Delete mock_torch

PiperOrigin-RevId: 515343503

Author: Etienne Pot

CHANGELOG.md

@@ -23,8 +23,7 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 ## [Unreleased]
 
-*   Add `torch` support (experimental). Require to call
-    `dca.activate_torch_support()`
+*   **Add `torch` support!**
 
 ## [1.3.0] - 2023-01-16
 

dataclass_array/__init__.py

@@ -37,9 +37,6 @@
 from dataclass_array.ops import stack
 from dataclass_array.vectorization import vectorize_method
 
-# TODO(epot): Remove once Torch has better numpy API
-from etils.enp import activate_torch_support
-
 # `dca.testing` do not depend on pytest or other heavy deps, so is safe to
 # import
 from dataclass_array import testing

dataclass_array/array_dataclass.py

@@ -541,17 +541,6 @@ def _get_xnp(f: _ArrayField) -> enp.NpModule:
       return None
     xnp = _infer_xnp(xnps)
 
-    if (
-        enp.lazy.has_torch
-        and xnp is enp.lazy.torch
-        and not hasattr(enp.lazy.torch, '__etils_np_mode__')
-    ):
-      raise ValueError(
-          'torch support currently require to call:\n'
-          'import dataclass_array as dca\n'
-          'dca.activate_torch_support()'
-      )
-
     def _cast_field(f: _ArrayField) -> None:
       try:
         new_value = np_utils.asarray(

dataclass_array/array_dataclass_test.py

@@ -28,7 +28,7 @@
 import tensorflow as tf
 
 # Activate the fixture
-enable_torch_tf_np_mode = enp.testing.enable_torch_tf_np_mode
+enable_tf_np_mode = enp.testing.set_tnp
 
 # TODO(epot): Test dtype `complex`, `str`
 

dataclass_array/import_test.py

@@ -25,7 +25,6 @@
 
 import dataclass_array as dca
 from etils import enp
-import pytest
 
 
 @dataclasses.dataclass(frozen=True)

dataclass_array/utils/np_utils_test.py

@@ -23,7 +23,7 @@
 import pytest
 
 # Activate the fixture
-enable_torch_tf_np_mode = enp.testing.enable_torch_tf_np_mode
+enable_tf_np_mode = enp.testing.set_tnp
 
 
 @enp.testing.parametrize_xnp()

dataclass_array/vectorization_test.py

@@ -34,7 +34,7 @@
 X1 = 5
 
 # Activate the fixture
-enable_torch_tf_np_mode = enp.testing.enable_torch_tf_np_mode
+enable_tf_np_mode = enp.testing.set_tnp
 
 
 @pytest.mark.parametrize(