Greatly limits how the dtype and device can be used in Tensor

Previously, setting `dtype` or `device` in the Tensor constructor
would incur a cast or copy operation respectively if the input data
were not already of the correct data type or on the right device.
These would be immediately evaluated, which could lead to extremely long tracing times.

Additionally, the implicit behavior is not in line with the rest of Tripy's
API design. With this change, users must explicitly call `cast` or `copy`.
This makes it clear *when* such an operation is actually happening.

Author: pranavm-nvidia

tripy/docs/README.md

@@ -73,7 +73,7 @@ def relu(input: "nvtripy.Tensor") -> "nvtripy.Tensor":
     .. code-block:: python
         :linenos:
 
-        input = tp.Tensor([1., 2., 3., 4.], dtype=tp.float32)
+        input = tp.Tensor([1., 2., 3., 4.])
         output = tp.relu(input)
 
         t = torch.tensor([1, 2, 3, 4], dtype=torch.float32) # doc: omit

tripy/docs/pre0_user_guides/01-compiler.md

@@ -74,7 +74,7 @@ fast_geglu = tp.compile(layer, args=[inp_info])
 out0 = fast_geglu(inp)
 
 # Now use an input with a different shape: (2, 2):
-inp1 = tp.Tensor([[1., 2.], [2., 3.]], dtype=tp.float32).eval()
+inp1 = tp.Tensor([[1., 2.], [2., 3.]]).eval()
 out1 = fast_geglu(inp1)
 ```
 

tripy/examples/nanogpt/example.py

@@ -132,7 +132,7 @@ def main():
         # Append sampled index to the running sequence and continue
         idx = torch.from_dlpack(idx).to(torch.int32)
         idx = torch.concat([idx, idx_next], dim=1).to(torch.int32)
-        idx = tp.Tensor(idx, device=tp.device("gpu"))
+        idx = tp.Tensor(idx)
 
     response = encoder.decode(torch.from_dlpack(idx[0, :]).tolist())
     end_time = time.perf_counter()

tripy/examples/segment-anything-model-v2/sam2/modeling/position_encoding.py

@@ -113,8 +113,7 @@ def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> No
         if scale is None or scale <= 0.0:
             scale = 1.0
         self.positional_encoding_gaussian_matrix = tp.Tensor(
-            scale * np.random.randn(2, num_pos_feats).astype(np.float32),
-            dtype=tp.float32,
+            scale * np.random.randn(2, num_pos_feats).astype(np.float32)
         )
 
     def _pe_encoding(self, coords: tp.Tensor) -> tp.Tensor:

tripy/examples/segment-anything-model-v2/sam2/modeling/sam/mask_decoder.py

@@ -332,7 +332,7 @@ def _get_stability_scores(self, mask_logits):
         stability_delta = self.dynamic_multimask_stability_delta
         area_i = tp.cast(tp.sum(tp.cast(mask_logits > stability_delta, tp.float32), dim=-1), self.dtype)
         area_u = tp.cast(tp.sum(tp.cast(mask_logits > -stability_delta, tp.float32), dim=-1), self.dtype)
-        stability_scores = tp.where(area_u > 0, area_i / area_u, tp.Tensor(1.0, dtype=self.dtype))
+        stability_scores = tp.where(area_u > 0, area_i / area_u, 1.0)
         return stability_scores
 
     def _dynamic_multimask_via_stability(self, all_mask_logits, all_iou_scores):

tripy/notebooks/resnet50.ipynb

@@ -460,11 +460,7 @@
     "\n",
     "    # Preprocess the image for inference\n",
     "    processed_image = processor(image, return_tensors=\"np\")[\"pixel_values\"]\n",
-    "    tp_image = tp.Tensor(\n",
-    "        processed_image,\n",
-    "        dtype=tp.float32,\n",
-    "        device=tp.device(\"gpu\"),\n",
-    "    )\n",
+    "    tp_image = tp.copy(tp.Tensor(processed_image), tp.device(\"gpu\"))\n",
     "\n",
     "    # Run inference with compiled model and get predicted label\n",
     "    label_idx = compiled_resnet_classifier(tp_image)\n",

tripy/nvtripy/frontend/dimension_size.py

@@ -34,7 +34,7 @@ def __init__(self, data: int, name: Optional[str] = None) -> None:
             data: The value of the DimensionSize, which should be a scalar integer.
             name: An optional name.
         """
-        super().__init__(data=data, dtype=int32, name=name)
+        super().__init__(data=data, name=name)
 
     def __int__(self) -> int:
         return self.tolist()

tripy/nvtripy/frontend/module/embedding.py

@@ -53,7 +53,7 @@ def __init__(self, num_embeddings: int, embedding_dim: int, dtype: datatype.dtyp
 
             embedding.weight = tp.iota(embedding.weight.shape)
 
-            input = tp.Tensor([0, 2], dtype=tp.int32)
+            input = tp.Tensor([0, 2])
             output = embedding(input)
 
             assert np.array_equal(cp.from_dlpack(output).get(), cp.from_dlpack(embedding.weight).get()[[0,2], :])

tripy/nvtripy/frontend/module/module.py

@@ -118,14 +118,14 @@ class Module:
         class AddBias(tp.Module):
             def __init__(self):
                 super().__init__()
-                self.bias = tp.Tensor([1.0, 1.0], dtype=tp.float32)
+                self.bias = tp.Tensor([1.0, 1.0])
 
             def forward(self, x):
                 return x + self.bias
 
         add_bias = AddBias()
 
-        input = tp.Tensor([1.0, 1.0], dtype=tp.float32)
+        input = tp.Tensor([1.0, 1.0])
         output = add_bias(input)
 
         assert np.array_equal(cp.from_dlpack(output).get(), np.array([2.0, 2.0]))

tripy/nvtripy/frontend/ops/cast.py

@@ -61,7 +61,7 @@ def cast(input: "nvtripy.Tensor", dtype: "nvtripy.dtype") -> "nvtripy.Tensor":
     .. code-block:: python
         :linenos:
 
-        input = tp.Tensor([1, 2], dtype=tp.int32)
+        input = tp.Tensor([1, 2])
         output = tp.cast(input, tp.float32)
 
         assert np.array_equal(cp.from_dlpack(output).get(), np.array([1, 2], dtype=np.float32))