Remove IREE helper functions (pybamm-team#5082)

* Remove most of IREE from the solvers

* Cleanup

* Remove most other jax and IREE code blocks

* Update changelog

* Remove some helper functions

Author: kratman

src/pybamm/__init__.py

@@ -1,8 +1,5 @@
 from pybamm.version import __version__
 
-# Demote expressions to 32-bit floats/ints - option used for IDAKLU-MLIR compilation
-demote_expressions_to_32bit = False
-
 # Utility classes and methods
 from .util import root_dir
 from .util import Timer, TimerTime, FuzzyDict

src/pybamm/expression_tree/operations/evaluate_python.py

@@ -603,54 +603,9 @@ def __init__(self, symbol: pybamm.Symbol):
             static_argnums=self._static_argnums,
         )
 
-    def _demote_constants(self):
-        """Demote 64-bit constants (f64, i64) to 32-bit (f32, i32)"""
-        if not pybamm.demote_expressions_to_32bit:
-            return  # pragma: no cover
-        self._constants = EvaluatorJax._demote_64_to_32(self._constants)
-
-    @classmethod
-    def _demote_64_to_32(cls, c):
-        """Demote 64-bit operations (f64, i64) to 32-bit (f32, i32)"""
-
-        if not pybamm.demote_expressions_to_32bit:
-            return c
-        if isinstance(c, float):
-            c = jax.numpy.float32(c)
-        if isinstance(c, int):
-            c = jax.numpy.int32(c)
-        if isinstance(c, np.int64):
-            c = c.astype(jax.numpy.int32)
-        if isinstance(c, np.ndarray):
-            if c.dtype == np.float64:
-                c = c.astype(jax.numpy.float32)
-            if c.dtype == np.int64:
-                c = c.astype(jax.numpy.int32)
-        if isinstance(c, jax.numpy.ndarray):
-            if c.dtype == jax.numpy.float64:
-                c = c.astype(jax.numpy.float32)
-            if c.dtype == jax.numpy.int64:
-                c = c.astype(jax.numpy.int32)
-        if isinstance(
-            c, pybamm.expression_tree.operations.evaluate_python.JaxCooMatrix
-        ):
-            if c.data.dtype == np.float64:
-                c.data = c.data.astype(jax.numpy.float32)
-            if c.row.dtype == np.int64:
-                c.row = c.row.astype(jax.numpy.int32)
-            if c.col.dtype == np.int64:
-                c.col = c.col.astype(jax.numpy.int32)
-        if isinstance(c, dict):
-            c = {key: EvaluatorJax._demote_64_to_32(value) for key, value in c.items()}
-        if isinstance(c, tuple):
-            c = tuple(EvaluatorJax._demote_64_to_32(value) for value in c)
-        if isinstance(c, list):
-            c = [EvaluatorJax._demote_64_to_32(value) for value in c]
-        return c
-
     @property
     def _constants(self):
-        return tuple(map(EvaluatorJax._demote_64_to_32, self.__constants))
+        return self.__constants
 
     @_constants.setter
     def _constants(self, value):

src/pybamm/solvers/idaklu_solver.py

@@ -526,13 +526,6 @@ def __setstate__(self, d):
             options=self._options,
         )
 
-    def _check_mlir_conversion(self, name, mlir: str):
-        if mlir.count("f64") > 0:  # pragma: no cover
-            warnings.warn(f"f64 found in {name} (x{mlir.count('f64')})", stacklevel=2)
-
-    def _demote_64_to_32(self, x: pybamm.EvaluatorJax):
-        return pybamm.EvaluatorJax._demote_64_to_32(x)
-
     @property
     def supports_parallel_solve(self):
         return True

src/pybamm/solvers/processed_variable_computed.py

@@ -140,11 +140,6 @@ def _unroll_nnz(self, realdata=None):
             nnz = sp.nnz()
             numel = sp.numel()
             row = sp.row()
-        elif "nnz" in dir(self.base_variables_casadi[0]):  # IREE fcn
-            sp = self.base_variables_casadi[0]
-            nnz = sp.nnz
-            numel = sp.numel
-            row = sp.row
         if nnz != numel:
             data = [None] * len(realdata)
             for datak in range(len(realdata)):

tests/unit/test_expression_tree/test_operations/test_evaluate_python.py

@@ -672,68 +672,6 @@ def test_evaluator_jax_inputs(self):
         result = evaluator(inputs={"a": 2})
         assert result == 4
 
-    @pytest.mark.skipif(not pybamm.has_jax(), reason="jax or jaxlib is not installed")
-    def test_evaluator_jax_demotion(self):
-        for demote in [True, False]:
-            pybamm.demote_expressions_to_32bit = demote  # global flag
-            target_dtype = "32" if demote else "64"
-            if demote:
-                # Test only works after conversion to jax.numpy
-                for c in [
-                    1.0,
-                    1,
-                ]:
-                    assert (
-                        str(pybamm.EvaluatorJax._demote_64_to_32(c).dtype)[-2:]
-                        == target_dtype
-                    )
-            for c in [
-                np.float64(1.0),
-                np.int64(1),
-                np.array([1.0], dtype=np.float64),
-                np.array([1], dtype=np.int64),
-                jax.numpy.array([1.0], dtype=np.float64),
-                jax.numpy.array([1], dtype=np.int64),
-            ]:
-                assert (
-                    str(pybamm.EvaluatorJax._demote_64_to_32(c).dtype)[-2:]
-                    == target_dtype
-                )
-            for c in [
-                {key: np.float64(1.0) for key in ["a", "b"]},
-            ]:
-                expr_demoted = pybamm.EvaluatorJax._demote_64_to_32(c)
-                assert all(
-                    str(c_v.dtype)[-2:] == target_dtype
-                    for c_k, c_v in expr_demoted.items()
-                )
-            for c in [
-                (np.float64(1.0), np.float64(2.0)),
-                [np.float64(1.0), np.float64(2.0)],
-            ]:
-                expr_demoted = pybamm.EvaluatorJax._demote_64_to_32(c)
-                assert all(str(c_i.dtype)[-2:] == target_dtype for c_i in expr_demoted)
-            for dtype in [
-                np.float64,
-                jax.numpy.float64,
-            ]:
-                c = pybamm.JaxCooMatrix([0, 1], [0, 1], dtype([1.0, 2.0]), (2, 2))
-                c_demoted = pybamm.EvaluatorJax._demote_64_to_32(c)
-                assert all(
-                    str(c_i.dtype)[-2:] == target_dtype for c_i in c_demoted.data
-                )
-            for dtype in [
-                np.int64,
-                jax.numpy.int64,
-            ]:
-                c = pybamm.JaxCooMatrix(
-                    dtype([0, 1]), dtype([0, 1]), [1.0, 2.0], (2, 2)
-                )
-                c_demoted = pybamm.EvaluatorJax._demote_64_to_32(c)
-                assert all(str(c_i.dtype)[-2:] == target_dtype for c_i in c_demoted.row)
-                assert all(str(c_i.dtype)[-2:] == target_dtype for c_i in c_demoted.col)
-            pybamm.demote_expressions_to_32bit = False
-
     @pytest.mark.skipif(not pybamm.has_jax(), reason="jax or jaxlib is not installed")
     def test_jax_coo_matrix(self):
         A = pybamm.JaxCooMatrix([0, 1], [0, 1], [1.0, 2.0], (2, 2))