torchax 0.0.4__py3-none-any.whl → 0.0.6__py3-none-any.whl

torchax/ops/mappings.py

@@ -7,7 +7,7 @@ import torch.utils.dlpack as torchdl
 import torch.utils._mode_utils as mode_utils
 
 
-def t2j(t):
+def t2j(t, use_dlpack=True):
   is_bool = False
   if t.dtype == torch.bool:
     is_bool = True
@@ -18,9 +18,14 @@ def t2j(t):
   if not t.is_contiguous():
     t = t.contiguous()
 
-  try:
-    res = jaxdl.from_dlpack(t)
-  except Exception:
+  res = None
+  if use_dlpack:
+    try:
+      res = jaxdl.from_dlpack(t)
+    except Exception:
+      pass
+
+  if res is None:
     # https://github.com/google/jax/issues/7657
     # https://github.com/google/jax/issues/17784
     if t.dtype == torch.bfloat16:
@@ -37,61 +42,98 @@ def t2j(t):
   return res
 
 
-def j2t(x):
+def j2t(x, use_dlpack=True):
   with mode_utils.no_dispatch(), torch._C.DisableTorchFunction():
-    try:
-      dl = jaxdl.to_dlpack(x)
-      res = torchdl.from_dlpack(dl)
-    except Exception:
+    res = None
+    if use_dlpack:
+      try:
+        dl = jaxdl.to_dlpack(x)
+        res = torchdl.from_dlpack(dl)
+      except Exception:
+        res = None
+
+    orig_dtype = None
+    if res is None:
+      orig_dtype = None
+      if x.dtype == jnp.bfloat16.dtype:
+        orig_dtype = x.dtype
+        x = x.astype(jnp.float32.dtype)
       res = torch.from_numpy(numpy.asarray(x))
+
     if x.dtype == jnp.bool_:
       res = res.to(torch.bool)
+
+    if orig_dtype is not None:
+      res = res.to(j2t_dtype(orig_dtype))
     return res
 
+
 TORCH_DTYPE_TO_JAX = {
     # NO_MAPPING        : jnp.float0.dtype (signless scalar int),
-    torch.bool          : jnp.bool_.dtype,
+    torch.bool:
+        jnp.bool_.dtype,
     # NO_MAPPING        : jnp.int4.dtype,
-    torch.int8          : jnp.int8.dtype,
-    torch.int16         : jnp.int16.dtype,
-    torch.int32         : jnp.int32.dtype,
-    torch.int64         : jnp.int64.dtype,
-    torch.long          : jnp.int64.dtype,
+    torch.int8:
+        jnp.int8.dtype,
+    torch.int16:
+        jnp.int16.dtype,
+    torch.int32:
+        jnp.int32.dtype,
+    torch.int64:
+        jnp.int64.dtype,
+    torch.long:
+        jnp.int64.dtype,
     # NO_MAPPING        : jnp.uint4
-    torch.uint8         : jnp.uint8.dtype,
-    torch.uint16        : jnp.uint16.dtype,
-    torch.uint32        : jnp.uint32.dtype,
-    torch.uint64        : jnp.uint64.dtype,
+    torch.uint8:
+        jnp.uint8.dtype,
+    torch.uint16:
+        jnp.uint16.dtype,
+    torch.uint32:
+        jnp.uint32.dtype,
+    torch.uint64:
+        jnp.uint64.dtype,
     # NO_MAPPING        : jnp.float8_e4m3b11fnuz.dtype,
-    torch.float8_e4m3fn : jnp.float8_e4m3fn.dtype,
+    torch.float8_e4m3fn:
+        jnp.float8_e4m3fn.dtype,
     # NO_MAPPING        : jnp.float8_e4m3fnuz.dtype,
-    torch.float8_e5m2   : jnp.float8_e5m2.dtype,
+    torch.float8_e5m2:
+        jnp.float8_e5m2.dtype,
     # NO_MAPPING        : jnp.float8_e5m2fnuz.dtype,
-    torch.bfloat16      : jnp.bfloat16.dtype,
-    torch.half          : jnp.float16.dtype,
-    torch.float16       : jnp.float16.dtype,
-    torch.float32       : jnp.float32.dtype,
-    torch.float64       : jnp.float64.dtype,
-    torch.double        : jnp.double.dtype,
-    torch.complex64     : jnp.complex64.dtype,
-    torch.complex128    : jnp.complex128.dtype,
-    None                : None,
+    torch.bfloat16:
+        jnp.bfloat16.dtype,
+    torch.half:
+        jnp.float16.dtype,
+    torch.float16:
+        jnp.float16.dtype,
+    torch.float32:
+        jnp.float32.dtype,
+    torch.float64:
+        jnp.float64.dtype,
+    torch.double:
+        jnp.double.dtype,
+    torch.complex64:
+        jnp.complex64.dtype,
+    torch.complex128:
+        jnp.complex128.dtype,
+    None:
+        None,
 }
 
-JAX_DTYPE_TO_TORCH = {
-  value: key for key, value in TORCH_DTYPE_TO_JAX.items()
-}
+JAX_DTYPE_TO_TORCH = {value: key for key, value in TORCH_DTYPE_TO_JAX.items()}
 # Add imprecise mappings for some JAX dtypes which don't have torch analogues
 JAX_DTYPE_TO_TORCH[jnp.dtype('int4')] = torch.int8
 JAX_DTYPE_TO_TORCH[jnp.dtype('uint4')] = torch.uint8
 
+
 def t2j_dtype(dtype):
   if dtype not in TORCH_DTYPE_TO_JAX:
-    raise RuntimeError(f'Attempting to convert unknown type: {dtype} to jax type,')
+    raise RuntimeError(
+        f'Attempting to convert unknown type: {dtype} to jax type,')
   return TORCH_DTYPE_TO_JAX[dtype]
 
 
 def j2t_dtype(dtype):
   if dtype not in JAX_DTYPE_TO_TORCH:
-    raise RuntimeError(f'Attempting to convert unknown type: {dtype} to torch type,')
+    raise RuntimeError(
+        f'Attempting to convert unknown type: {dtype} to torch type,')
   return JAX_DTYPE_TO_TORCH[dtype]

torchax/ops/op_base.py

@@ -4,6 +4,7 @@ import jax.numpy as jnp
 import numpy as np
 import torch
 from torchax.ops import mappings
+from torchax.view import View
 from torchax import types
 import sys
 
@@ -12,31 +13,55 @@ from typing import Callable, Optional, ParamSpec, Concatenate
 
 class InplaceOp:
 
-    def __init__(self, functional_op, replace=False, position_to_mutate=0):
-        self.functional = functional_op
-        self.replace = replace
-        self.position_to_mutate = position_to_mutate
-
-    def __call__(self, *args, **kwargs):
-        to_mutate = args[0]
-        if self.replace:
-          to_mutate._elem = self.functional(*args, **kwargs)._elem
-        else:
-          to_mutate.copy_(self.functional(*args, **kwargs))
-        return to_mutate
+  def __init__(self,
+               functional_op,
+               replace=False,
+               position_to_mutate=0,
+               is_jax_func=False):
+    self.functional = functional_op
+    self.replace = replace
+    self.position_to_mutate = position_to_mutate
+    self.is_jax_func = is_jax_func
+
+  def __call__(self, *args, **kwargs):
+    to_mutate = args[self.position_to_mutate]
+    view_value = to_mutate
+    if isinstance(to_mutate, View):
+      view_value = to_mutate.torch()
+      # Convert the target View to a Tensor, and
+      # leave the rest args as is. If other args are
+      # also View, they will be converted to tensors
+      # in the self.functional dispatch.
+    env = view_value._env
+    if self.is_jax_func:
+      view_value, args, kwargs = env.t2j_iso((view_value, args, kwargs))
+      new_value_jax = self.functional(view_value, *args[1:], **kwargs)
+      new_value = env.j2t_iso(new_value_jax)
+    else:
+      new_value = self.functional(view_value, *args[1:], **kwargs)
+
+    if isinstance(to_mutate, View):
+      to_mutate.update(new_value)
+    else:
+      if self.replace:
+        to_mutate._elem = new_value._elem
+      else:
+        to_mutate.copy_(new_value)
+    return to_mutate
 
 
 class OutVariant:
 
-    def __call__(self, *args, **kwargs):
-        to_mutate = kwargs['out']
-        del kwargs['out']
-        to_mutate._elem = self.functional(*args, **kwargs)._elem
-        return to_mutate
-
+  def __call__(self, *args, **kwargs):
+    to_mutate = kwargs['out']
+    del kwargs['out']
+    to_mutate._elem = self.functional(*args, **kwargs)._elem
+    return to_mutate
 
 
 P = ParamSpec('P')
+
+
 def convert_dtype(use_default_dtype: bool = True):
   """Converts `dtype` kwarg of function from torch to JAX.
 
@@ -48,6 +73,7 @@ def convert_dtype(use_default_dtype: bool = True):
   """
 
   def decorator(func: types.TorchCallable):
+
     @functools.wraps(func)
     def wrapper(*args: P.args,
                 dtype: Optional[torch.dtype] = None,
@@ -66,7 +92,8 @@ def convert_dtype(use_default_dtype: bool = True):
   return decorator
 
 
-def maybe_convert_constant_dtype(val: Optional[types.JaxValue], dtype: Optional[jnp.dtype]):
+def maybe_convert_constant_dtype(val: Optional[types.JaxValue],
+                                 dtype: Optional[jnp.dtype]):
   """Optionally converts scalar constant's dtype using `numpy`
 
   Use in cases where you require a constant and can't handle a traced array.
@@ -81,24 +108,24 @@ def maybe_convert_constant_dtype(val: Optional[types.JaxValue], dtype: Optional[
 
 
 def promote_int_input(f: Callable[Concatenate[jax.Array, P], types.JaxValue]):
-   """If the first argument is an int array, promote it to float32."""
-   @functools.wraps(f)
-   def wrapper(x: jax.Array, *args: P.args, **kwargs: P.kwargs):
-      if x.dtype in [jnp.int8, jnp.int16, jnp.int32, jnp.int64]:
-        x = x.astype(mappings.t2j_dtype(torch.get_default_dtype()))
+  """If the first argument is an int array, promote it to float32."""
+
+  @functools.wraps(f)
+  def wrapper(x: jax.Array, *args: P.args, **kwargs: P.kwargs):
+    if x.dtype in [jnp.int8, jnp.int16, jnp.int32, jnp.int64]:
+      x = x.astype(mappings.t2j_dtype(torch.get_default_dtype()))
 
-      return f(x, *args, **kwargs)
+    return f(x, *args, **kwargs)
 
-   return wrapper
+  return wrapper
 
 
-def foreach_loop(
-  seq: jax.Array, fn: Callable[[jax.Array, jax.Array], jax.Array], init_val=0.0
-):
+def foreach_loop(seq: jax.Array,
+                 fn: Callable[[jax.Array, jax.Array], jax.Array],
+                 init_val=0.0):
   """Run `fn` for each element of 1D array `seq`.
 
   Similar to `functools.reduce`, but implemented with `jax.lax.fori_loop`."""
   assert len(seq.shape) == 1
-  return jax.lax.fori_loop(
-    0, len(seq), lambda i, carry: fn(carry, seq[i]), init_val
-  )
+  return jax.lax.fori_loop(0, len(seq), lambda i, carry: fn(carry, seq[i]),
+                           init_val)

torchax/ops/ops_registry.py

@@ -7,44 +7,49 @@ from typing import Union, Dict
 
 @dataclasses.dataclass
 class Operator:
-    torch_op: TorchCallable
-    func: Union[TorchCallable, JaxCallable]
-    is_jax_function: bool
-    is_user_defined: bool
-    needs_env: bool
+  torch_op: TorchCallable
+  func: Union[TorchCallable, JaxCallable]
+  is_jax_function: bool
+  is_user_defined: bool
+  needs_env: bool
+  is_view_op: bool
 
 
 all_aten_ops: Dict[TorchCallable, Operator] = {}
 all_torch_functions: Dict[TorchCallable, Operator] = {}
 
 
-def register_torch_dispatch_op(
-    aten_op, impl_callable,
-    is_jax_function=True,
-    is_user_defined=False,
-    needs_env=False,
-):
-    op = Operator(
-        aten_op, impl_callable,
-        is_jax_function=is_jax_function,
-        is_user_defined=is_user_defined,
-        needs_env=needs_env)
-    if aten_op in all_aten_ops:
-        logging.warning(f'Duplicate op registration for {aten_op}')
-    all_aten_ops[aten_op] = op
-    return impl_callable
-
-
-def register_torch_function_op(
-    torch_func, impl_callable,
-    is_jax_function=True,
-    is_user_defined=False,
-    needs_env=False,
-):
-    op = Operator(
-        torch_func, impl_callable,
-        is_jax_function=is_jax_function,
-        is_user_defined=is_user_defined,
-        needs_env=needs_env)
-    all_torch_functions[torch_func] = op
-    return impl_callable
+def register_torch_dispatch_op(aten_op,
+                               impl_callable,
+                               is_jax_function=True,
+                               is_user_defined=False,
+                               needs_env=False,
+                               is_view_op=False):
+  op = Operator(
+      aten_op,
+      impl_callable,
+      is_jax_function=is_jax_function,
+      is_user_defined=is_user_defined,
+      needs_env=needs_env,
+      is_view_op=is_view_op)
+  if aten_op in all_aten_ops:
+    logging.warning(f'Duplicate op registration for {aten_op}')
+  all_aten_ops[aten_op] = op
+  return impl_callable
+
+
+def register_torch_function_op(torch_func,
+                               impl_callable,
+                               is_jax_function=True,
+                               is_user_defined=False,
+                               needs_env=False,
+                               is_view_op=False):
+  op = Operator(
+      torch_func,
+      impl_callable,
+      is_jax_function=is_jax_function,
+      is_user_defined=is_user_defined,
+      needs_env=needs_env,
+      is_view_op=is_view_op)
+  all_torch_functions[torch_func] = op
+  return impl_callable