Misc fixes: (#9491)

Author: qihqi

torchax/test/test_interop.py

@@ -5,6 +5,7 @@
 from torchax import interop, jax_device
 import torchax
 import jax
+import jax.numpy as jnp
 
 
 def is_tpu_available():
@@ -171,6 +172,17 @@ def test_to_jax_device(self):
         self.assertEqual(d.jax_device.platform, "cpu")
         self.assertEqual(d.device.type, "jax")
 
+  def test_torch_jax_view_dtype(self):
+    dtype = torch.float32
+    self.assertEqual(interop.jax_view(dtype), jnp.float32.dtype)
+    self.assertEqual(interop.torch_view(interop.jax_view(dtype)), dtype)
+    dtype = torch.bfloat16
+    self.assertEqual(interop.jax_view(dtype), jnp.bfloat16.dtype)
+    self.assertEqual(interop.torch_view(interop.jax_view(dtype)), dtype)
+    dtype = torch.int32
+    self.assertEqual(interop.jax_view(dtype), jnp.int32.dtype)
+    self.assertEqual(interop.torch_view(interop.jax_view(dtype)), dtype)
+
 
 if __name__ == '__main__':
   unittest.main()

torchax/torchax/__init__.py

@@ -49,10 +49,11 @@ def extract_jax(mod: torch.nn.Module, env=None):
   states = env.t2j_copy(states)
 
   #@jax.jit
-  def jax_func(states, inputs):
-    (states, inputs) = env.j2t_iso((states, inputs))
+  def jax_func(states, args, kwargs=None):
+    (states, args, kwargs) = env.j2t_iso((states, args, kwargs))
     with env:
-      res = torch.func.functional_call(mod, states, inputs, tie_weights=False)
+      res = torch.func.functional_call(
+          mod, states, args, kwargs, tie_weights=False)
     return env.t2j_iso(res)
 
   return states, jax_func

torchax/torchax/device_module.py

@@ -1,3 +1,6 @@
+import torch
+
+
 def _is_in_bad_fork():
   return False
 
@@ -24,3 +27,7 @@ def is_available():
 
 def current_device():
   return 0
+
+
+def get_amp_supported_dtype():
+  return [torch.float16, torch.bfloat16]

torchax/torchax/interop.py

@@ -11,6 +11,7 @@
 from jax.experimental.shard_map import shard_map
 from torchax import tensor
 from torchax import util
+from torchax.ops import mappings
 import torchax
 
 from torchax.types import JaxValue, TorchValue, JaxCallable, TorchCallable
@@ -183,8 +184,8 @@ def _torch_view(t: JaxValue) -> TorchValue:
   if isinstance(t, jax.Array):
     # TODO
     return tensor.Tensor(t, torchax.default_env())
-  if isinstance(t, type(jnp.int32)):
-    return tensor.t2j_type(t)
+  if isinstance(t, jnp.dtype):
+    return mappings.j2t_dtype(t)
   if callable(t):  # t is a JaxCallable
     return functools.partial(call_jax, t)
   # regular types are not changed
@@ -201,7 +202,7 @@ def _jax_view(t: TorchValue) -> JaxValue:
     assert isinstance(t, tensor.Tensor) or isinstance(t, tensor.View), type(t)
     return t.jax()
   if isinstance(t, type(torch.int32)):
-    return tensor.t2j_dtype(t)
+    return mappings.t2j_dtype(t)
 
   # torch.nn.Module needs special handling
   if not isinstance(t, torch.nn.Module) and callable(t):  # t is a TorchCallable