Allow tensor proxy to handle kwargs (#80)

### What change is being made

Fix todo.

### Why this change is being made

Missing feature.

### Tested

Added unit test.

GitOrigin-RevId: 7f7acb063249b5491dfd7ec2c10c6b47ee56832e

Author: dbellicoso-bdai

python/exploy/exporter/core/tensor_proxy.py

@@ -159,12 +159,13 @@ def __torch_function__(cls, func, types, args=(), kwargs=None):
         # Convert all `TensorProxy`` elements in `args` to `torch.Tensor` objects.
         new_args = args_to_tensor(args=args)
 
-        # todo: Convert TensorProxy to Tensor in kwargs.
+        # Convert TensorProxy to Tensor in kwargs.
+        new_kwargs = {k: args_to_tensor(args=(v,))[0] for k, v in kwargs.items()}
 
         return (
-            func.__wrapped__(*new_args, **kwargs)
+            func.__wrapped__(*new_args, **new_kwargs)
             if hasattr(func, "__wrapped__")
-            else func(*new_args, **kwargs)
+            else func(*new_args, **new_kwargs)
         )
 
     def __repr__(self):

python/exploy/exporter/core/tests/test_tensor_proxy.py

@@ -366,6 +366,17 @@ def test_torch_cat(self):
         expected = torch.cat([tensor1, tensor2], dim=0)
         assert torch.equal(result, expected)
 
+    def test_torch_cat_with_kwargs(self):
+        """Test using TensorProxy with torch.cat using keyword arguments."""
+        tensor1 = torch.rand((2, 3))
+        tensor2 = torch.rand((2, 3))
+        proxy1 = TensorProxy(tensor1, split_dim=0)
+        proxy2 = TensorProxy(tensor2, split_dim=0)
+
+        result = torch.cat(tensors=[proxy1, proxy2], dim=0)
+        expected = torch.cat([tensor1, tensor2], dim=0)
+        assert torch.equal(result, expected)
+
     def test_repr(self):
         """Test string representation."""
         tensor = torch.rand((2, 3, 4))