[WIP] Add Numpy Backend

keras_core/backend/__init__.py

@@ -36,5 +36,8 @@
 elif backend() == "torch":
     print_msg("Using PyTorch backend.")
     from keras_core.backend.torch import *  # noqa: F403
+elif backend() == "numpy":
+    print_msg("Using NumPy backend.")
+    from keras_core.backend.numpy import *  # noqa: F403
 else:
     raise ValueError(f"Unable to import backend : {backend()}")

keras_core/backend/numpy/__init__.py

@@ -0,0 +1,15 @@
+from keras_core.backend.numpy import core
+from keras_core.backend.numpy import math
+from keras_core.backend.numpy import nn
+from keras_core.backend.numpy import numpy
+from keras_core.backend.numpy import random
+from keras_core.backend.numpy.core import DYNAMIC_SHAPES_OK
+from keras_core.backend.numpy.core import Variable
+from keras_core.backend.numpy.core import cast
+from keras_core.backend.numpy.core import compute_output_spec
+from keras_core.backend.numpy.core import cond
+from keras_core.backend.numpy.core import convert_to_tensor
+from keras_core.backend.numpy.core import is_tensor
+from keras_core.backend.numpy.core import name_scope
+from keras_core.backend.numpy.core import shape
+from keras_core.backend.numpy.core import vectorized_map

keras_core/backend/numpy/core.py

@@ -0,0 +1,89 @@
+from contextlib import nullcontext
+
+import numpy as np
+from tensorflow import nest
+
+from keras_core.backend.common import KerasVariable
+from keras_core.backend.common import standardize_dtype
+from keras_core.backend.common.keras_tensor import KerasTensor
+from keras_core.backend.common.stateless_scope import StatelessScope
+
+DYNAMIC_SHAPES_OK = False
+
+
+class Variable(KerasVariable):
+    def _initialize(self, value):
+        self._value = np.array(value, dtype=self._dtype)
+
+    def _direct_assign(self, value):
+        self._value = value
+
+    def _convert_to_tensor(self, value, dtype=None):
+        return convert_to_tensor(value, dtype=dtype)
+
+    # Overload native accessor.
+    def __array__(self):
+        return self
+
+
+def convert_to_tensor(x, dtype=None):
+    if dtype is not None:
+        dtype = standardize_dtype(dtype)
+    if isinstance(x, Variable):
+        if dtype and dtype != x.dtype:
+            return x.value.astype(dtype)
+        return x.value
+    return np.array(x, dtype=dtype)
+
+
+def is_tensor(x):
+    if isinstance(x, np.ndarray):
+        return True
+    return False
+
+
+def shape(x):
+    # This will work as long as we disallow
+    # dynamic shapes in NumPy.
+    return x.shape
+
+
+def cast(x, dtype):
+    return convert_to_tensor(x, dtype=dtype)
+
+
+def cond(pred, true_fn, false_fn):
+    if pred:
+        return true_fn
+    return false_fn
+
+
+def name_scope(name):
+    # There is no need for a named context for NumPy.
+    return nullcontext()
+
+
+def vectorized_map(function, elements):
+    return np.vectorize(function)(elements)
+
+
+# Shape / dtype inference util
+def compute_output_spec(fn, *args, **kwargs):
+    with StatelessScope():
+
+        def convert_keras_tensor_to_numpy(x):
+            if isinstance(x, KerasTensor):
+                return np.ones(x.shape, dtype=x.dtype)
+            return x
+
+        args, kwargs = nest.map_structure(
+            convert_keras_tensor_to_numpy, (args, kwargs)
+        )
+        np_out = fn(*args, **kwargs)
+
+        def convert_numpy_to_keras_tensor(x):
+            if isinstance(x, np.ndarray):
+                return KerasTensor(x.shape, x.dtype)
+            return x
+
+        return nest.map_structure(convert_numpy_to_keras_tensor, np_out)

keras_core/backend/numpy/layer.py

@@ -0,0 +1,2 @@
+class NumpyLayer:
+    pass

keras_core/backend/numpy/math.py

@@ -0,0 +1,51 @@
+import numpy as np
+
+
+def segment_sum(data, segment_ids, num_segments=None, sorted=False):
+    if num_segments is None:
+        # Segment ids must be less than the number of segments.
+        num_segments = np.amax(segment_ids) + 1
+
+    if sorted:
+        result = np.zeros(num_segments, dtype=data.dtype)
+        np.add.at(result, segment_ids, data)
+    else:
+        sort_indices = np.argsort(segment_ids)
+        sorted_segment_ids = segment_ids[sort_indices]
+        sorted_data = data[sort_indices]
+
+        result = np.zeros(num_segments, dtype=data.dtype)
+        np.add.at(result, sorted_segment_ids, sorted_data)
+
+    return result
+
+
+def top_k(x, k, sorted=False):
+    sorted_indices = np.argsort(x, axis=-1)[..., ::-1]
+    sorted_values = np.sort(x, axis=-1)[..., ::-1]
+
+    if sorted:
+        # Take the k largest values.
+        top_k_values = sorted_values[..., :k]
+        top_k_indices = sorted_indices[..., :k]
+    else:
+        # Partition the array such that all values larger than the k-th
+        # largest value are to the right of it.
+        top_k_values = np.partition(x, -k, axis=-1)[..., -k:]
+        top_k_indices = np.argpartition(x, -k, axis=-1)[..., -k:]
+
+        # Get the indices in sorted order.
+        idx = np.argsort(-top_k_values, axis=-1)
+
+        # Get the top k values and their indices.
+        top_k_values = np.take_along_axis(top_k_values, idx, axis=-1)
+        top_k_indices = np.take_along_axis(top_k_indices, idx, axis=-1)
+
+    return top_k_values, top_k_indices
+
+
+def in_top_k(targets, predictions, k):
+    top_k_indices = top_k(predictions, k, sorted=True)[1]
+    targets = targets[..., np.newaxis]
+    mask = targets == top_k_indices
+    return np.any(mask, axis=-1)