lunny/keras
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
1d265f4b67
keras/keras_core/operations/numpy.py
Francois Chollet bb1d4eeb18 Add layer testing infra.
2023-04-22 18:03:15 -07:00

2960 lines
82 KiB
Python
Raw Blame History

"""
MANIFEST:
abs
absolute
add
all
amax
amin
append
arange
arccos
arcsin
arctan
arctan2
argmax
argmin
argsort
array
average
broadcast_to
ceil
clip
concatenate
conj
conjugate
copy
cos
count_nonzero
cross
cumprod
cumsum
diag
diagonal
diff
divide
dot
dtype
einsum
empty
equal
exp
expand_dims
expm1
eye
flip
floor
full
full_like
greater
greater_equal
hstack
identity
imag
interp
isclose
isfinite
isinf
isnan
less
less_equal
linspace
log
log10
log1p
log2
logaddexp
logical_and
logical_not
logical_or
logspace
matmul
max
maximum
mean
median
meshgrid
mgrid
min
minimum
mod
moveaxis
multiply
nan_to_num
ndim
nonzero
not_equal
ones
ones_like
outer
pad
percentile
power
prod
ravel
real
reciprocal
repeat
reshape
roll
round
sign
sin
size
sort
split
sqrt
square
squeeze
stack
std
subtract
sum
swapaxes
take
take_along_axis
tan
tensordot
tile
trace
transpose
tri
tril
triu
true_divide
vdot
vstack
where
zeros
zeros_like
"""
import numpy as np
from keras_core import backend
from keras_core.backend import KerasTensor
from keras_core.backend import any_symbolic_tensors
from keras_core.operations.operation import Operation
def broadcast_shapes(shape1, shape2):
# Broadcast input shapes to a unified shape.
# Convert to list for mutability.
shape1 = list(shape1)
shape2 = list(shape2)
origin_shape1 = shape1
origin_shape2 = shape2
if len(shape1) > len(shape2):
shape2 = [None] * (len(shape1) - len(shape2)) + shape2
if len(shape1) < len(shape2):
shape1 = [None] * (len(shape2) - len(shape1)) + shape1
output_shape = list(shape1)
for i in range(len(shape1)):
if shape1[i] == 1:
output_shape[i] = shape2[i]
elif shape1[i] is None:
output_shape[i] = shape2[i]
else:
if shape2[i] == 1 or shape2[i] is None or shape2[i] == shape1[i]:
output_shape[i] = shape1[i]
else:
raise ValueError(
"Cannot broadcast shape, the failure dim has value "
f"{shape1[i]}, which cannot be broadcasted to {shape2[i]}. "
f"Input shapes are: {origin_shape1} and {origin_shape2}."
)
return output_shape
def reduce_shape(shape, axis=None, keepdims=False):
shape = list(shape)
if axis is None:
if keepdims:
output_shape = [1 for _ in range(shape)]
else:
output_shape = []
return output_shape
if keepdims:
for ax in axis:
shape[ax] = 1
return shape
else:
for ax in axis:
shape[ax] = -1
output_shape = list(filter((-1).__ne__, shape))
return output_shape
def shape_equal(shape1, shape2, axis=None, allow_none=True):
"""Check if two shapes are equal.
Args:
shape1: A tuple or list of integers.
shape2: A tuple or list of integers.
axis: int or list/tuple of ints, defaults to None. If specified, the
shape check will ignore the axes specified by `axis`.
allow_none: bool, defaults to True. If True, None in the shape will
match any value.
"""
if len(shape1) != len(shape2):
return False
shape1 = list(shape1)
shape2 = list(shape2)
if axis is not None:
for ax in axis:
shape1[ax] = -1
shape2[ax] = -1
if allow_none:
for i in range(len(shape1)):
if shape1[i] is None:
shape1[i] = shape2[i]
if shape2[i] is None:
shape2[i] = shape1[i]
return shape1 == shape2
class Absolute(Operation):
def call(self, x):
return backend.execute("absolute", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def absolute(x):
if any_symbolic_tensors((x,)):
return Absolute().symbolic_call(x)
return backend.execute("absolute", x)
class Abs(Absolute):
pass
def abs(x):
return absolute(x)
class Add(Operation):
def call(self, x1, x2):
return backend.numpy.add(x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
output_shape = broadcast_shapes(x1_shape, x2_shape)
return KerasTensor(output_shape, dtype=x1.dtype)
def add(x1, x2):
if any_symbolic_tensors((x1, x2)):
return Add().symbolic_call(x1, x2)
return backend.numpy.add(x1, x2)
class All(Operation):
def __init__(self, axis=None, keepdims=False):
super().__init__()
if isinstance(axis, int):
self.axis = [axis]
else:
self.axis = axis
self.keepdims = keepdims
def call(self, x):
return backend.execute(
"all",
x,
axis=self.axis,
keepdims=self.keepdims,
)
def compute_output_spec(self, x):
return KerasTensor(
reduce_shape(
x.shape,
axis=self.axis,
keepdims=self.keepdims,
),
dtype=x.dtype,
)
def all(x, axis=None, keepdims=False):
if any_symbolic_tensors((x,)):
return All(axis=axis, keepdims=keepdims).symbolic_call(x)
return backend.execute("all", x, axis=axis, keepdims=keepdims)
class Amax(Operation):
def __init__(self, axis=None, keepdims=False):
super().__init__()
if isinstance(axis, int):
axis = [axis]
self.axis = axis
self.keepdims = keepdims
def call(self, x):
return backend.execute(
"amax",
x,
axis=self.axis,
keepdims=self.keepdims,
)
def compute_output_spec(self, x):
return KerasTensor(
reduce_shape(x.shape, axis=self.axis, keepdims=self.keepdims),
dtype=x.dtype,
)
def amax(x, axis=None, keepdims=False):
if any_symbolic_tensors((x,)):
return All(axis=axis, keepdims=keepdims).symbolic_call(x)
return backend.execute("amax", x, axis=axis, keepdims=keepdims)
class Amin(Operation):
def __init__(self, axis=None, keepdims=False):
super().__init__()
if isinstance(axis, int):
axis = [axis]
self.axis = axis
self.keepdims = keepdims
def call(self, x):
return backend.execute(
"amin", x, axis=self.axis, keepdims=self.keepdims
)
def compute_output_spec(self, x):
return KerasTensor(
reduce_shape(x.shape, axis=self.axis, keepdims=self.keepdims),
dtype=x.dtype,
)
def amin(x, axis=None, keepdims=False):
if any_symbolic_tensors((x,)):
return All(axis=axis, keepdims=keepdims).symbolic_call(x)
return backend.execute("amin", x, axis=axis, keepdims=keepdims)
class Append(Operation):
def __init__(self, axis=None):
super().__init__()
self.axis = axis
def call(self, x1, x2):
return backend.execute("append", x1, x2, axis=self.axis)
def compute_output_spec(self, x1, x2):
x1_shape = x1.shape
x2_shape = x2.shape
if self.axis is None:
if None in x1_shape or None in x2_shape:
output_shape = [None]
else:
output_shape = [int(np.prod(x1_shape) + np.prod(x2_shape))]
return KerasTensor(output_shape, dtype=x1.dtype)
if not shape_equal(x1_shape, x2_shape, [self.axis]):
raise ValueError(
"`append` requires inputs to have the same shape except the "
f"`axis={self.axis}`, but received shape {x1_shape} and "
f"{x2_shape}."
)
output_shape = list(x1_shape)
output_shape[self.axis] = x1_shape[self.axis] + x2_shape[self.axis]
return KerasTensor(output_shape, dtype=x1.dtype)
def append(
x1,
x2,
axis=None,
):
if any_symbolic_tensors((x1, x2)):
return Append(axis=axis).symbolic_call(x1, x2)
return backend.execute("append", x1, x2, axis=axis)
class Arange(Operation):
def call(self, start, stop=None, step=None, dtype=None):
if stop is None:
start, stop = 0, start
if step is None:
step = 1
return backend.execute("arange", start, stop, step=step, dtype=dtype)
def compute_output_spec(self, start, stop=None, step=None, dtype=None):
if stop is None:
start, stop = 0, start
if step is None:
step = 1
output_shape = [np.ceil((stop - start) / step).astype(int)]
return KerasTensor(output_shape, dtype=dtype)
def arange(start, stop=None, step=None, dtype=None):
if stop is None:
start, stop = 0, start
if step is None:
step = 1
return backend.execute("arange", start, stop, step=step, dtype=dtype)
class Arccos(Operation):
def call(self, x):
return backend.execute("arccos", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def arccos(x):
if any_symbolic_tensors((x,)):
return Arccos().symbolic_call(x)
return backend.execute("arccos", x)
class Arcsin(Operation):
def call(self, x):
return backend.execute("arcsin", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def arcsin(x):
if any_symbolic_tensors((x,)):
return Arcsin().symbolic_call(x)
return backend.execute("arcsin", x)
class Arctan(Operation):
def call(self, x):
return backend.execute("arctan", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def arctan(x):
if any_symbolic_tensors((x,)):
return Arctan().symbolic_call(x)
return backend.execute("arctan", x)
class Arctan2(Operation):
def call(self, x1, x2):
return backend.execute("arctan2", x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
outputs_shape = broadcast_shapes(x1_shape, x2_shape)
return KerasTensor(outputs_shape, dtype=x1.dtype)
def arctan2(x1, x2):
if any_symbolic_tensors((x1, x2)):
return Arctan2().symbolic_call(x1, x2)
return backend.execute("arctan2", x1, x2)
class Argmax(Operation):
def __init__(self, axis=None):
super().__init__()
self.axis = axis
def call(self, x):
return backend.execute("argmax", x, axis=self.axis)
def compute_output_spec(self, x):
if self.axis is None:
return KerasTensor([], dtype="int32")
return KerasTensor(
reduce_shape(x.shape, axis=[self.axis]), dtype="int32"
)
def argmax(x, axis=None):
if any_symbolic_tensors((x,)):
return Argmax(axis=axis).symbolic_call(x)
return backend.execute("argmax", x, axis=axis)
class Argmin(Operation):
def __init__(self, axis=None):
super().__init__()
self.axis = axis
def call(self, x):
return backend.execute("argmin", x, axis=self.axis)
def compute_output_spec(self, x):
if self.axis is None:
return KerasTensor([], dtype="int32")
return KerasTensor(
reduce_shape(x.shape, axis=[self.axis]), dtype="int32"
)
def argmin(x, axis=None):
if any_symbolic_tensors((x,)):
return Argmin(axis=axis).symbolic_call(x)
return backend.execute("argmin", x, axis=axis)
class Argsort(Operation):
def __init__(self, axis=-1):
super().__init__()
self.axis = axis
def call(self, x):
return backend.execute("argsort", x, axis=self.axis)
def compute_output_spec(self, x):
if self.axis is None:
return KerasTensor([int(np.prod(x.shape))], dtype="int32")
return KerasTensor(x.shape, dtype="int32")
def argsort(x, axis=-1):
if any_symbolic_tensors((x,)):
return Argsort(axis=axis).symbolic_call(x)
return backend.execute("argsort", x, axis=axis)
class Array(Operation):
def call(self, x, dtype=None):
return backend.execute("array", x, dtype=dtype)
def compute_output_spec(self, x, dtype=None):
return KerasTensor(x.shape, dtype=dtype)
def array(x, dtype=None):
if any_symbolic_tensors((x,)):
return Array().symbolic_call(x, dtype=dtype)
return backend.execute("array", x, dtype=dtype)
class Average(Operation):
def __init__(self, axis=None):
super().__init__()
# np.average() does not support axis as tuple as declared by the
# docstring, it only supports int or None.
self.axis = axis
def call(self, x, weights=None):
return backend.execute("average", x, weights=weights, axis=self.axis)
def compute_output_spec(self, x, weights=None):
if weights is not None:
shape_match = shape_equal(x.shape, weights.shape, allow_none=True)
if self.axis is not None:
shape_match_on_axis = shape_equal(
[x.shape[self.axis]], weights.shape, allow_none=True
)
if self.axis is None:
if weights is None or shape_match:
return KerasTensor(
[],
dtype=x.dtype,
)
else:
raise ValueError(
"`weights` must have the same shape as `x` when "
f"`axis=None`, but received `weights.shape={weights.shape}`"
f" and `x.shape={x.shape}`."
)
if weights is None or shape_match_on_axis or shape_match:
return KerasTensor(
reduce_shape(x.shape, axis=[self.axis]),
dtype=x.dtype,
)
else:
# `weights` can either be a 1D array of length `x.shape[axis]` or
# of the same shape as `x`.
raise ValueError(
"`weights` must have the same size as `x` at "
f"`axis={self.axis}` but received "
f"`weights.shape={weights.shape}` while x.shape at "
f"`{self.axis}` is `{x.shape[self.axis]}`."
)
def average(x, axis=None, weights=None):
if any_symbolic_tensors((x,)):
return Average(axis=axis).symbolic_call(x, weights=weights)
return backend.execute("average", x, weights=weights, axis=axis)
class BroadcastTo(Operation):
def __init__(self, shape):
super().__init__()
self.shape = shape
def call(self, x):
return backend.execute("broadcast_to", x, self.shape)
def compute_output_spec(self, x):
# Catch broadcasting errors for clear error messages.
broadcast_shapes(x.shape, self.shape)
return KerasTensor(self.shape, dtype=x.dtype)
def broadcast_to(x, shape):
if any_symbolic_tensors((x,)):
return BroadcastTo(shape=shape).symbolic_call(x)
return backend.execute("broadcast_to", x, shape)
class Ceil(Operation):
def call(self, x):
return backend.execute("ceil", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def ceil(x):
if any_symbolic_tensors((x,)):
return Ceil().symbolic_call(x)
return backend.execute("ceil", x)
class Clip(Operation):
def __init__(self, x_min, x_max):
super().__init__()
self.x_min = x_min
self.x_max = x_max
def call(self, x):
return backend.execute("clip", x, self.x_min, self.x_max)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def clip(x, x_min, x_max):
if any_symbolic_tensors((x,)):
return Clip(x_min, x_max).symbolic_call(x)
return backend.execute("clip", x, x_min, x_max)
class Concatenate(Operation):
def __init__(self, axis=0):
super().__init__()
if axis is None:
raise ValueError("`axis` cannot be None for `concatenate`.")
self.axis = axis
def call(self, xs):
return backend.execute("concatenate", xs, axis=self.axis)
def compute_output_spec(self, xs):
first_shape = xs[0].shape
total_size_on_axis = 0
for x in xs:
if not shape_equal(
x.shape, first_shape, axis=[self.axis], allow_none=True
):
raise ValueError(
"Every value in `xs` must have the same shape except on "
f"the `axis` dim. But found element of shape {x.shape}, "
f"which is different from the first element's "
f"shape {first_shape}."
)
if total_size_on_axis is None or x.shape[self.axis] is None:
total_size_on_axis = None
else:
total_size_on_axis += x.shape[self.axis]
output_shape = list(first_shape)
output_shape[self.axis] = total_size_on_axis
return KerasTensor(output_shape, dtype=x.dtype)
def concatenate(xs, axis=0):
if any_symbolic_tensors(xs):
return Concatenate(axis=axis).symbolic_call(xs)
return backend.execute("concatenate", xs, axis=axis)
class Conjugate(Operation):
def call(self, x):
return backend.execute("conjugate", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def conjugate(x):
if any_symbolic_tensors((x,)):
return Conjugate().symbolic_call(x)
return backend.execute("conjugate", x)
class Conj(Conjugate):
pass
def conj(x):
return conjugate(x)
class Copy(Operation):
def call(self, x):
return backend.execute("copy", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def copy(x):
if any_symbolic_tensors((x,)):
return Copy().symbolic_call(x)
return backend.execute("copy", x)
class Cos(Operation):
def call(self, x):
return backend.execute("cos", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def cos(x):
if any_symbolic_tensors((x,)):
return Cos().symbolic_call(x)
return backend.execute("cos", x)
class CountNonzero(Operation):
def __init__(self, axis=None):
super().__init__()
if isinstance(axis, int):
self.axis = (axis,)
else:
self.axis = axis
def call(self, x):
return backend.execute("count_nonzero", x, axis=self.axis)
def compute_output_spec(self, x):
return KerasTensor(
reduce_shape(x.shape, axis=self.axis),
dtype="int32",
)
def count_nonzero(x, axis=None):
if any_symbolic_tensors((x,)):
return CountNonzero(axis=axis).symbolic_call(x)
return backend.execute("count_nonzero", x, axis=axis)
class Cross(Operation):
def __init__(self, axisa=-1, axisb=-1, axisc=-1, axis=None):
super().__init__()
if axis is not None:
self.axisa = axis
self.axisb = axis
self.axisc = axis
else:
self.axisa = axisa
self.axisb = axisb
self.axisc = axisc
def call(self, x1, x2):
return backend.execute(
"cross", x1, x2, self.axisa, self.axisb, self.axisc
)
def compute_output_spec(self, x1, x2):
x1_shape = list(x1.shape)
x2_shape = list(x2.shape)
x1_value_size = x1_shape[self.axisa]
x2_value_size = x2_shape[self.axisa]
del x1_shape[self.axisa]
del x2_shape[self.axisb]
output_shape = broadcast_shapes(x1_shape, x2_shape)
if x1_value_size is not None and x1_value_size not in (2, 3):
raise ValueError(
"`x1`'s dim on `axis={axisa}` must be either 2 or 3, but "
f"received: {x1_value_size}"
)
if x2_value_size is not None and x2_value_size not in (2, 3):
raise ValueError(
"`x2`'s dim on `axis={axisb}` must be either 2 or 3, but "
f"received: {x2_value_size}"
)
if x1_value_size == 3 or x2_value_size == 3:
value_size = [3]
else:
value_size = []
output_shape = (
output_shape[: self.axisc] + value_size + output_shape[self.axisc :]
)
return KerasTensor(output_shape, dtype=x1.dtype)
def cross(x1, x2, axisa=-1, axisb=-1, axisc=-1, axis=None):
if any_symbolic_tensors((x1, x2)):
return Cross(
axisa=axisa, axisb=axisb, axisc=axisc, axis=axis
).symbolic_call(x1, x2)
return backend.execute(
"cross",
x1,
x2,
axisa=axisa,
axisb=axisb,
axisc=axisc,
axis=axis,
)
class Cumprod(Operation):
def __init__(self, axis=None):
super().__init__()
self.axis = axis
def call(self, x):
return backend.execute("cumprod", x, axis=self.axis)
def compute_output_spec(self, x):
if self.axis is None:
if None in x.shape:
output_shape = (None,)
else:
output_shape = (int(np.prod(x.shape)),)
return KerasTensor(output_shape, dtype="int32")
return KerasTensor(x.shape, dtype=x.dtype)
def cumprod(x, axis=None):
if any_symbolic_tensors((x,)):
return Cumprod(axis=axis).symbolic_call(x)
return backend.execute("cumprod", x, axis=axis)
class Cumsum(Operation):
def __init__(self, axis=None):
super().__init__()
self.axis = axis
def call(self, x):
return backend.execute("cumsum", x, axis=self.axis)
def compute_output_spec(self, x):
if self.axis is None:
if None in x.shape:
output_shape = (None,)
else:
output_shape = (int(np.prod(x.shape)),)
return KerasTensor(output_shape, dtype="int32")
return KerasTensor(x.shape, dtype=x.dtype)
def cumsum(x, axis=None):
if any_symbolic_tensors((x,)):
return Cumsum(axis=axis).symbolic_call(x)
return backend.execute("cumsum", x, axis=axis)
class Diag(Operation):
def __init__(self, k=0):
super().__init__()
self.k = k
def call(self, x):
return backend.execute("diag", x, k=self.k)
def compute_output_spec(self, x):
x_shape = x.shape
if len(x_shape) == 1:
if x_shape[0] is None:
output_shape = [None, None]
else:
output_shape = [
x_shape[0] + int(np.abs(self.k)),
x_shape[0] + int(np.abs(self.k)),
]
elif len(x_shape) == 2:
if None in x_shape:
output_shape = [None]
else:
shorter_side = np.minimum(x_shape[0], x_shape[1])
if self.k > 0:
remaining = x_shape[1] - self.k
else:
remaining = x_shape[0] + self.k
output_shape = [
int(np.maximum(0, np.minimum(remaining, shorter_side)))
]
else:
raise ValueError(
f"`x` must be 1-D or 2-D, but received shape {x.shape}."
)
return KerasTensor(output_shape, dtype=x.dtype)
def diag(x, k=0):
if any_symbolic_tensors((x,)):
return Diag(k=k).symbolic_call(x)
return backend.execute("diag", x, k=k)
class Diagonal(Operation):
def __init__(self, offset=0, axis1=0, axis2=1):
super().__init__()
self.offset = offset
self.axis1 = axis1
self.axis2 = axis2
def call(self, x):
return backend.execute(
"diagonal",
x,
offset=self.offset,
axis1=self.axis1,
axis2=self.axis2,
)
def compute_output_spec(self, x):
x_shape = list(x.shape)
if len(x_shape) < 2:
raise ValueError(
"`diagonal` requires an array of at least two dimensions, but "
"`x` is of shape {x.shape}."
)
shape_2d = [x_shape[self.axis1], x_shape[self.axis2]]
x_shape[self.axis1] = -1
x_shape[self.axis2] = -1
output_shape = list(filter((-1).__ne__, x_shape))
if None in shape_2d:
diag_shape = [None]
else:
shorter_side = np.minimum(shape_2d[0], shape_2d[1])
if self.offset > 0:
remaining = shape_2d[1] - self.offset
else:
remaining = shape_2d[0] + self.offset
diag_shape = [
int(np.maximum(0, np.minimum(remaining, shorter_side)))
]
output_shape = output_shape + diag_shape
return KerasTensor(output_shape, dtype=x.dtype)
def diagonal(x, offset=0, axis1=0, axis2=1):
if any_symbolic_tensors((x,)):
return Diagonal(
offset=offset,
axis1=axis1,
axis2=axis2,
).symbolic_call(x)
return backend.execute(
"diagonal",
x,
offset=offset,
axis1=axis1,
axis2=axis2,
)
class Dot(Operation):
def call(self, x1, x2):
return backend.execute("dot", x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
if x1_shape == [] or x2_shape == []:
return multiply(x1, x2)
if len(x1_shape) == 1 and len(x2_shape) == 1:
return KerasTensor([], dtype=x1.dtype)
if len(x2_shape) == 1:
if x1_shape[-1] != x2_shape[0]:
raise ValueError(
"Shape must match on the last axis of `x1` and `x2` when "
"`x1` is N-d array while `x2` is 1-D, but receive shape "
f"`x1.shape={x1.shape}` and x2.shape=`{x2.shape}`."
)
return KerasTensor(x1_shape[:-1], dtype=x1.dtype)
if x1_shape[-1] != x2_shape[-2]:
raise ValueError(
"Shape must match on the last axis of `x1` and second last "
"axis of `x2` when `x1` is N-d array while `x2` is M-D, but "
f"received `x1.shape={x1.shape}` and x2.shape=`{x2.shape}`."
)
del x1_shape[-1]
del x2_shape[-2]
return KerasTensor(x1_shape + x2_shape, dtype=x1.dtype)
class Empty(Operation):
def call(self, shape, dtype="float32"):
return backend.execute("empty", shape, dtype=dtype)
def compute_output_spec(self, shape, dtype="float32"):
return KerasTensor(shape, dtype=dtype)
def empty(shape, dtype="float32"):
return backend.execute("empty", shape, dtype=dtype)
class Equal(Operation):
def call(self, x1, x2):
return backend.execute("equal", x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
output_shape = broadcast_shapes(x1_shape, x2_shape)
return KerasTensor(output_shape, dtype=x1.dtype)
def equal(x1, x2):
if any_symbolic_tensors((x1, x2)):
return Equal().symbolic_call(x1, x2)
return backend.execute("equal", x1, x2)
class Exp(Operation):
def call(self, x):
return backend.execute("exp", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def exp(x):
if any_symbolic_tensors((x,)):
return Exp().symbolic_call(x)
return backend.execute("exp", x)
class ExpandDims(Operation):
def __init__(self, axis):
super().__init__()
if isinstance(axis, list):
raise ValueError(
"The `axis` argument to `expand_dims` should be an integer, "
f"but received a list: {axis}."
)
self.axis = axis
def call(self, x):
return backend.execute("expand_dims", x, self.axis)
def compute_output_spec(self, x):
x_shape = list(x.shape)
if self.axis < 0:
axis = len(x.shape) + 1 + self.axis
else:
axis = self.axis
output_shape = x_shape[:axis] + [1] + x_shape[axis:]
return KerasTensor(output_shape, dtype=x.dtype)
def expand_dims(x, axis):
if any_symbolic_tensors((x,)):
return ExpandDims(axis=axis).symbolic_call(x)
return backend.execute("expand_dims", x, axis)
class Expm1(Operation):
def call(self, x):
return backend.execute("expm1", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def expm1(x):
if any_symbolic_tensors((x,)):
return Expm1().symbolic_call(x)
return backend.execute("expm1", x)
class Flip(Operation):
def __init__(self, axis=None):
super().__init__()
self.axis = axis
def call(self, x):
return backend.execute("flip", x, axis=self.axis)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def flip(x, axis=None):
if any_symbolic_tensors((x,)):
return Flip(axis=axis).symbolic_call(x)
return backend.execute("flip", x, axis=axis)
class Floor(Operation):
def call(self, x):
return backend.execute("floor", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def floor(x):
if any_symbolic_tensors((x,)):
return Floor().symbolic_call(x)
return backend.execute("floor", x)
class Full(Operation):
def call(self, shape, fill_value, dtype=None):
return backend.execute("full", shape, fill_value, dtype=dtype)
def compute_output_spec(self, shape, fill_value, dtype=None):
return KerasTensor(shape, dtype=dtype)
def full(shape, fill_value, dtype=None):
return backend.execute("full", shape, fill_value, dtype=dtype)
class FullLike(Operation):
def call(self, x, fill_value, dtype=None):
return backend.execute("full_like", x, fill_value, dtype=dtype)
def compute_output_spec(self, x, fill_value, dtype=None):
return KerasTensor(x.shape, dtype=dtype)
def full_like(x, fill_value, dtype=None):
if any_symbolic_tensors((x,)):
return FullLike().symbolic_call(x, fill_value, dtype=dtype)
return backend.execute("full_like", x, fill_value, dtype=dtype)
class GetItem(Operation):
def call(self, x, key):
if not isinstance(key, int):
# TODO: support slicing.
raise ValueError(
"Only scalar int keys are supported at this time. Cannot "
f"process key {key}"
)
return x[key]
def compute_output_spec(self, x, key):
if not isinstance(key, int):
# TODO: support slicing.
raise ValueError(
"Only scalar int keys are supported at this time. Cannot "
f"process key {key}"
)
if len(x.shape) == 0:
raise ValueError(
f"Too many indices for array: array is scalar "
f"but index {key} was requested. A scalar array "
"cannot be indexed."
)
if x.shape[0] is not None and key >= x.shape[0]:
raise ValueError(
f"Array has shape {x.shape} "
f"but out-of-bound index {key} was requested."
)
return KerasTensor(x.shape[1:], dtype=x.dtype)
def get_item(x, key):
if any_symbolic_tensors((x,)):
return GetItem().symbolic_call(x, key)
if not isinstance(key, int):
# TODO: support slicing.
raise ValueError(
"Only scalar int keys are supported at this time. Cannot "
f"process key {key}"
)
return x[key]
class Greater(Operation):
def call(self, x1, x2):
return backend.execute("greater", x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
output_shape = broadcast_shapes(x1_shape, x2_shape)
return KerasTensor(output_shape, dtype=x1.dtype)
def greater(x1, x2):
if any_symbolic_tensors((x1, x2)):
return Greater().symbolic_call(x1, x2)
return backend.execute("greater", x1, x2)
class GreaterEqual(Operation):
def call(self, x1, x2):
return backend.execute("greater_equal", x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
output_shape = broadcast_shapes(x1_shape, x2_shape)
return KerasTensor(output_shape, dtype=x1.dtype)
def greater_equal(x1, x2):
if any_symbolic_tensors((x1, x2)):
return GreaterEqual().symbolic_call(x1, x2)
return backend.execute("greater_equal", x1, x2)
class Hstack(Operation):
def call(self, xs):
return backend.execute("hstack", xs)
def compute_output_spec(self, xs):
first_shape = xs[0].shape
total_size_on_axis = 0
for x in xs:
if not shape_equal(x.shape, first_shape, axis=[1], allow_none=True):
raise ValueError(
"Every value in `xs` must have the same shape except on "
f"the `axis` dim. But found element of shape {x.shape}, "
f"which is different from the first element's "
f"shape {first_shape}."
)
if total_size_on_axis is None or x.shape[1] is None:
total_size_on_axis = None
else:
total_size_on_axis += x.shape[1]
output_shape = list(first_shape)
output_shape[1] = total_size_on_axis
return KerasTensor(output_shape)
def hstack(xs):
if any_symbolic_tensors((xs,)):
return Hstack().symbolic_call(xs)
return backend.execute("hstack", xs)
class Identity(Operation):
def call(self, n, dtype="float32"):
return backend.execute("identity", n, dtype=dtype)
def compute_output_spec(self, n, dtype="float32"):
return KerasTensor([n, n], dtype=dtype)
def identity(n, dtype="float32"):
return backend.execute("identity", n, dtype=dtype)
class Imag(Operation):
def call(self, x):
return backend.execute("imag", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def imag(x):
if any_symbolic_tensors((x,)):
return Imag().symbolic_call(x)
return backend.execute("imag", x)
class Isclose(Operation):
def call(self, x1, x2):
return backend.execute("isclose", x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
output_shape = broadcast_shapes(x1_shape, x2_shape)
return KerasTensor(output_shape, dtype=x1.dtype)
def isclose(x1, x2):
if any_symbolic_tensors((x1, x2)):
return Isclose().symbolic_call(x1, x2)
return backend.execute("isclose", x1, x2)
class Isfinite(Operation):
def call(self, x):
return backend.execute("isfinite", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype="bool")
def isfinite(x):
if any_symbolic_tensors((x,)):
return Isfinite().symbolic_call(x)
return backend.execute("isfinite", x)
class Isinf(Operation):
def call(self, x):
return backend.execute("isinf", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype="bool")
def isinf(x):
if any_symbolic_tensors((x,)):
return Isinf().symbolic_call(x)
return backend.execute("isinf", x)
class Isnan(Operation):
def call(self, x):
return backend.execute("isnan", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype="bool")
def isnan(x):
if any_symbolic_tensors((x,)):
return Isnan().symbolic_call(x)
return backend.execute("isnan", x)
class Less(Operation):
def call(self, x1, x2):
return backend.execute("less", x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
output_shape = broadcast_shapes(x1_shape, x2_shape)
return KerasTensor(output_shape, dtype=x1.dtype)
def less(x1, x2):
if any_symbolic_tensors((x1, x2)):
return Less().symbolic_call(x1, x2)
return backend.execute("less", x1, x2)
class LessEqual(Operation):
def call(self, x1, x2):
return backend.execute("less_equal", x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
output_shape = broadcast_shapes(x1_shape, x2_shape)
return KerasTensor(output_shape, dtype=x1.dtype)
def less_equal(x1, x2):
if any_symbolic_tensors((x1, x2)):
return LessEqual().symbolic_call(x1, x2)
return backend.execute("less_equal", x1, x2)
class Linspace(Operation):
def __init__(
self, num=50, endpoint=True, retstep=False, dtype=float, axis=0
):
super().__init__()
self.num = num
self.endpoint = endpoint
self.retstep = retstep
self.dtype = dtype
self.axis = axis
def call(self, start, stop):
return backend.execute(
"linspace",
start,
stop,
num=self.num,
endpoint=self.endpoint,
retstep=self.retstep,
dtype=self.dtype,
axis=self.axis,
)
def compute_output_spec(self, start, stop):
start_shape = getattr(start, "shape", [])
stop_shape = getattr(stop, "shape", [])
output_shape = broadcast_shapes(start_shape, stop_shape)
if self.axis == -1:
output_shape = output_shape + [self.num]
elif self.axis >= 0:
output_shape = (
output_shape[: self.axis]
+ [self.num]
+ output_shape[self.axis :]
)
else:
output_shape = (
output_shape[: self.axis + 1]
+ [self.num]
+ output_shape[self.axis + 1 :]
)
dtype = self.dtype if self.dtype is not None else start.dtype
if self.retstep:
return (KerasTensor(output_shape, dtype=dtype), None)
return KerasTensor(output_shape, dtype=dtype)
def linspace(
start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis=0
):
if any_symbolic_tensors((start, stop)):
return Linspace(num, endpoint, retstep, dtype, axis)(start, stop)
return backend.execute(
"linspace",
start,
stop,
num=num,
endpoint=endpoint,
retstep=retstep,
dtype=dtype,
axis=axis,
)
class Log(Operation):
def call(self, x):
return backend.execute("log", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def log(x):
if any_symbolic_tensors((x,)):
return Log().symbolic_call(x)
return backend.execute("log", x)
class Log10(Operation):
def call(self, x):
return backend.execute("log10", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def log10(x):
if any_symbolic_tensors((x,)):
return Log10().symbolic_call(x)
return backend.execute("log10", x)
class Log1p(Operation):
def call(self, x):
return backend.execute("log1p", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def log1p(x):
if any_symbolic_tensors((x,)):
return Log10().symbolic_call(x)
return backend.execute("log1p", x)
class Log2(Operation):
def call(self, x):
return backend.execute("log2", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def log2(x):
if any_symbolic_tensors((x,)):
return Log2().symbolic_call(x)
return backend.execute("log2", x)
class Logaddexp(Operation):
def call(self, x1, x2):
return backend.execute("logaddexp", x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
output_shape = broadcast_shapes(x1_shape, x2_shape)
return KerasTensor(output_shape, dtype=x1.dtype)
def logaddexp(x1, x2):
if any_symbolic_tensors((x1, x2)):
return Logaddexp().symbolic_call(x1, x2)
return backend.execute("logaddexp", x1, x2)
class LogicalAnd(Operation):
def call(self, x1, x2):
return backend.execute("logical_and", x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
output_shape = broadcast_shapes(x1_shape, x2_shape)
return KerasTensor(output_shape, dtype=x1.dtype)
def logical_and(x1, x2):
if any_symbolic_tensors((x1, x2)):
return LogicalAnd().symbolic_call(x1, x2)
return backend.execute("logical_and", x1, x2)
class LogicalNot(Operation):
def call(self, x):
return backend.execute("logical_not", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def logical_not(x):
if any_symbolic_tensors((x,)):
return LogicalNot().symbolic_call(x)
return backend.execute("logical_not", x)
class LogicalOr(Operation):
def call(self, x1, x2):
return backend.execute("logical_or", x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
output_shape = broadcast_shapes(x1_shape, x2_shape)
return KerasTensor(output_shape, dtype=x1.dtype)
def logical_or(x1, x2):
if any_symbolic_tensors((x1, x2)):
return LogicalOr().symbolic_call(x1, x2)
return backend.execute("logical_or", x1, x2)
class Logspace(Operation):
def __init__(self, num=50, endpoint=True, base=10, dtype=float, axis=0):
super().__init__()
self.num = num
self.endpoint = endpoint
self.base = base
self.dtype = dtype
self.axis = axis
def call(self, start, stop):
return backend.execute(
"logspace",
start,
stop,
num=self.num,
endpoint=self.endpoint,
base=self.base,
dtype=self.dtype,
axis=self.axis,
)
def compute_output_spec(self, start, stop):
start_shape = getattr(start, "shape", [])
stop_shape = getattr(stop, "shape", [])
output_shape = broadcast_shapes(start_shape, stop_shape)
if self.axis == -1:
output_shape = output_shape + [self.num]
elif self.axis >= 0:
output_shape = (
output_shape[: self.axis]
+ [self.num]
+ output_shape[self.axis :]
)
else:
output_shape = (
output_shape[: self.axis + 1]
+ [self.num]
+ output_shape[self.axis + 1 :]
)
dtype = self.dtype if self.dtype is not None else start.dtype
return KerasTensor(output_shape, dtype=dtype)
def logspace(start, stop, num=50, endpoint=True, base=10, dtype=None, axis=0):
if any_symbolic_tensors((start, stop)):
return Logspace(num, endpoint, base, dtype, axis)(start, stop)
return backend.execute(
"logspace",
start,
stop,
num=num,
endpoint=endpoint,
base=base,
dtype=dtype,
axis=axis,
)
class Matmul(Operation):
def call(self, x1, x2):
return backend.numpy.matmul(x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
if len(x1_shape) == 1:
x1_shape = (1, x1_shape[0])
if len(x2_shape) == 1:
x2_shape = (x2_shape[0], 1)
if (
x1_shape[-1] is not None
and x2_shape[-2] is not None
and x1_shape[-1] != x2_shape[-2]
):
raise ValueError(
"Inner dimensions (`x1.shape[-1]` and `x2.shape[-2]`) must be "
f"equal, but received `x1.shape={x1.shape}` and "
f"`x2.shape={x2.shape}`."
)
leading_shape = broadcast_shapes(x1_shape[:-2], x2_shape[:-2])
last_2_dims_shape = [x1_shape[-2], x2_shape[-1]]
output_shape = leading_shape + last_2_dims_shape
if len(x1.shape) == 1:
del output_shape[-2]
if len(x2.shape) == 1:
del output_shape[-1]
return KerasTensor(output_shape, dtype=x1.dtype)
def matmul(x1, x2):
if any_symbolic_tensors((x1, x2)):
return Matmul().symbolic_call(x1, x2)
# The below conversion works around an outstanding JAX bug.
x1 = backend.convert_to_tensor(x1)
x2 = backend.convert_to_tensor(x2)
return backend.numpy.matmul(x1, x2)
class Max(Operation):
def __init__(self, axis=None, keepdims=False):
super().__init__()
if isinstance(axis, int):
self.axis = [axis]
else:
self.axis = axis
self.keepdims = keepdims
def call(self, x):
return backend.numpy.max(x, axis=self.axis, keepdims=self.keepdims)
def compute_output_spec(self, x):
return KerasTensor(
reduce_shape(x.shape, axis=self.axis, keepdims=self.keepdims),
dtype=x.dtype,
)
def max(x, axis=None, keepdims=False):
if any_symbolic_tensors((x,)):
return Max(axis=axis, keepdims=keepdims).symbolic_call(x)
return backend.numpy.max(x, axis=axis, keepdims=keepdims)
class Maximum(Operation):
def call(self, x1, x2):
return backend.execute("maximum", x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
output_shape = broadcast_shapes(x1_shape, x2_shape)
return KerasTensor(output_shape, dtype=x1.dtype)
def maximum(x1, x2):
if any_symbolic_tensors((x1, x2)):
return Maximum().symbolic_call(x1, x2)
return backend.execute("maximum", x1, x2)
class Meshgrid(Operation):
def __init__(self, indexing="xy"):
super().__init__()
if indexing not in ("xy", "ij"):
raise ValueError(
"Valid values for `indexing` are 'xy' and 'ij', "
"but received {index}."
)
self.indexing = indexing
def call(self, *x):
return backend.execute("meshgrid", *x, indexing=self.indexing)
def compute_output_spec(self, *x):
output_shape = []
for xi in x:
if len(xi.shape) == 0:
size = 1
else:
if None in xi.shape:
size = None
else:
size = int(np.prod(xi.shape))
output_shape.append(size)
if self.indexing == "ij":
return [KerasTensor(output_shape) for _ in range(len(x))]
tmp = output_shape[0]
output_shape[0] = output_shape[1]
output_shape[1] = tmp
return [KerasTensor(output_shape) for _ in range(len(x))]
def meshgrid(*x, indexing="xy"):
if any_symbolic_tensors(x):
return Meshgrid(indexing=indexing).symbolic_call(*x)
return backend.execute("meshgrid", *x, indexing=indexing)
class Min(Operation):
def __init__(self, axis=None, keepdims=False):
if isinstance(axis, int):
self.axis = [axis]
else:
self.axis = axis
self.keepdims = keepdims
def call(self, x):
return backend.execute("min", x, axis=self.axis, keepdims=self.keepdims)
def compute_output_spec(self, x):
return KerasTensor(
reduce_shape(x.shape, axis=self.axis, keepdims=self.keepdims),
dtype=x.dtype,
)
def min(x, axis=None, keepdims=False):
if any_symbolic_tensors((x,)):
return Min(axis=axis, keepdims=keepdims).symbolic_call(x)
return backend.execute("min", x, axis=axis, keepdims=keepdims)
class Minimum(Operation):
def call(self, x1, x2):
return backend.execute("minimum", x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
output_shape = broadcast_shapes(x1_shape, x2_shape)
return KerasTensor(output_shape, dtype=x1.dtype)
def minimum(x1, x2):
if any_symbolic_tensors((x1, x2)):
return Minimum().symbolic_call(x1, x2)
return backend.execute("minimum", x1, x2)
class Mod(Operation):
def call(self, x1, x2):
return backend.execute("mod", x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
output_shape = broadcast_shapes(x1_shape, x2_shape)
return KerasTensor(output_shape, dtype=x1.dtype)
def mod(x1, x2):
if any_symbolic_tensors((x1, x2)):
return Mod().symbolic_call(x1, x2)
return backend.execute("mod", x1, x2)
class Moveaxis(Operation):
def __init__(self, source, destination):
super().__init__()
if isinstance(source, int):
self.source = [source]
else:
self.source = source
if isinstance(destination, int):
self.destination = [destination]
else:
self.destination = destination
if len(self.source) != len(self.destination):
raise ValueError(
"`source` and `destination` arguments must have the same "
f"number of elements, but received `source={source}` and "
f"`destination={destination}`."
)
def call(self, x):
return backend.execute("moveaxis", x, self.source, self.destination)
def compute_output_spec(self, x):
x_shape = list(x.shape)
output_shape = [-1 for _ in range(len(x.shape))]
for sc, dst in zip(self.source, self.destination):
output_shape[dst] = x_shape[sc]
x_shape[sc] = -1
i, j = 0, 0
while i < len(output_shape):
while i < len(output_shape) and output_shape[i] != -1:
# Find the first dim unset.
i += 1
while j < len(output_shape) and x_shape[j] == -1:
# Find the first dim not being passed.
j += 1
if i == len(output_shape):
break
output_shape[i] = x_shape[j]
i += 1
j += 1
return KerasTensor(output_shape, dtype=x.dtype)
def moveaxis(x, source, destination):
if any_symbolic_tensors((x,)):
return Moveaxis(source, destination).symbolic_call(x)
return backend.execute(
"moveaxis", x, source=source, destination=destination
)
class Ndim(Operation):
def call(self, x):
return backend.execute(
"ndim",
x,
)
def compute_output_spec(self, x):
return KerasTensor([len(x.shape)])
def ndim(x):
if any_symbolic_tensors((x,)):
return Ndim().symbolic_call(x)
return backend.execute("ndim", x)
class Nonzero(Operation):
def call(self, x):
return backend.execute("nonzero", x)
def nonzero(x):
return backend.execute("nonzero", x)
class NotEqual(Operation):
def call(self, x1, x2):
return backend.execute("not_equal", x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
output_shape = broadcast_shapes(x1_shape, x2_shape)
return KerasTensor(output_shape, dtype=x1.dtype)
def not_equal(x1, x2):
if any_symbolic_tensors((x1, x2)):
return NotEqual().symbolic_call(x1, x2)
return backend.execute("not_equal", x1, x2)
class OnesLike(Operation):
def call(self, x, dtype=None):
return backend.execute("ones_like", x, dtype=dtype)
def compute_output_spec(self, x, dtype=None):
return KerasTensor(x.shape, dtype=dtype)
def ones_like(x, dtype=None):
if any_symbolic_tensors((x,)):
return OnesLike().symbolic_call(x, dtype=dtype)
return backend.execute("ones_like", x, dtype=dtype)
class Outer(Operation):
def call(self, x1, x2):
return backend.execute("outer", x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [1])
x2_shape = getattr(x2, "shape", [1])
if None in x1_shape:
x1_flatten_shape = None
else:
x1_flatten_shape = int(np.prod(x1_shape))
if None in x2_shape:
x2_flatten_shape = None
else:
x2_flatten_shape = int(np.prod(x2_shape))
output_shape = [x1_flatten_shape, x2_flatten_shape]
return KerasTensor(output_shape, dtype=x1.dtype)
def outer(x1, x2):
if any_symbolic_tensors((x1, x2)):
return Outer().symbolic_call(x1, x2)
return backend.execute("outer", x1, x2)
class Pad(Operation):
def __init__(self, pad_width, mode="constant"):
super().__init__()
self.pad_width = self._process_pad_width(pad_width)
self.mode = mode
def _process_pad_width(self, pad_width):
if isinstance(pad_width, int):
return ((pad_width, pad_width),)
if isinstance(pad_width, (tuple, list)) and isinstance(
pad_width[0], int
):
return (pad_width,)
first_len = len(pad_width[0])
for i, pw in enumerate(pad_width):
if len(pw) != first_len:
raise ValueError(
"`pad_width` should be a list of tuples of length 2 or "
f"1, but received {pad_width}."
)
if len(pw) == 1:
pad_width[i] = (pw[0], pw[0])
return pad_width
def call(self, x):
return backend.execute(
"pad", x, pad_width=self.pad_width, mode=self.mode
)
def compute_output_spec(self, x):
output_shape = list(x.shape)
if len(self.pad_width) == 1:
pad_width = [self.pad_width[0] for _ in range(len(output_shape))]
elif len(self.pad_width) == len(output_shape):
pad_width = self.pad_width
else:
raise ValueError(
"`pad_width` must have the same length as `x.shape`, but "
f"received {len(self.pad_width)} and {len(x.shape)}."
)
for i in range(len(output_shape)):
if output_shape[i] is None:
output_shape[i] = None
else:
output_shape[i] += pad_width[i][0] + pad_width[i][1]
return KerasTensor(output_shape, dtype=x.dtype)
def pad(x, pad_width, mode="constant"):
if any_symbolic_tensors((x,)):
return Pad(pad_width, mode=mode).symbolic_call(x)
return backend.execute("pad", x, pad_width, mode=mode)
class Prod(Operation):
def __init__(self, axis=None, keepdims=False, dtype=None):
super().__init__()
if isinstance(axis, int):
self.axis = [axis]
else:
self.axis = axis
self.keepdims = keepdims
self.dtype = dtype
def call(self, x):
return backend.execute(
"prod",
x,
axis=self.axis,
keepdims=self.keepdims,
dtype=self.dtype,
)
def compute_output_spec(self, x):
return KerasTensor(
reduce_shape(x.shape, axis=self.axis, keepdims=self.keepdims),
dtype=self.dtype,
)
def prod(x, axis=None, keepdims=False, dtype=None):
if any_symbolic_tensors((x,)):
return Prod(axis=axis, keepdims=keepdims, dtype=dtype).symbolic_call(x)
return backend.execute("prod", x, axis=axis, keepdims=keepdims, dtype=dtype)
class Ravel(Operation):
def call(self, x):
return backend.execute("ravel", x)
def compute_output_spec(self, x):
if None in x.shape:
output_shape = [
None,
]
else:
output_shape = [int(np.prod(x.shape))]
return KerasTensor(output_shape, dtype=x.dtype)
def ravel(x):
if any_symbolic_tensors((x,)):
return Ravel().symbolic_call(x)
return backend.execute("ravel", x)
class Real(Operation):
def call(self, x):
return backend.execute("real", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape)
def real(x):
if any_symbolic_tensors((x,)):
return Real().symbolic_call(x)
return backend.execute("real", x)
class Reciprocal(Operation):
def call(self, x):
return backend.execute("reciprocal", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape)
def reciprocal(x):
if any_symbolic_tensors((x,)):
return Reciprocal().symbolic_call(x)
return backend.execute("reciprocal", x)
class Repeat(Operation):
def __init__(self, repeats, axis=None):
super().__init__()
self.axis = axis
self.repeats = repeats
def call(self, x):
return backend.execute("repeat", x, self.repeats, axis=self.axis)
def compute_output_spec(self, x):
x_shape = list(x.shape)
if self.axis is None:
if None in x_shape:
return KerasTensor([None], dtype=x.dtype)
x_flatten_size = int(np.prod(x_shape))
if isinstance(self.repeats, int):
output_shape = [x_flatten_size * self.repeats]
else:
output_shape = [int(np.sum(self.repeats))]
return KerasTensor(output_shape, dtype=x.dtype)
size_on_ax = x_shape[self.axis]
output_shape = x_shape
if isinstance(self.repeats, int):
output_shape[self.axis] = size_on_ax * self.repeats
else:
output_shape[self.axis] = int(np.sum(self.repeats))
return KerasTensor(output_shape, dtype=x.dtype)
def repeat(x, repeats, axis=None):
if any_symbolic_tensors((x,)):
return Repeat(repeats, axis=axis).symbolic_call(x)
return backend.execute("repeat", x, repeats, axis=axis)
class Reshape(Operation):
def __init__(self, new_shape):
super().__init__()
self.new_shape = new_shape
def call(self, x):
return backend.execute("reshape", x, self.new_shape)
def compute_output_spec(self, x):
return KerasTensor(self.new_shape, dtype=x.dtype)
def reshape(x, new_shape):
if any_symbolic_tensors((x,)):
return Reshape(new_shape).symbolic_call(x)
return backend.execute("reshape", x, new_shape)
class Roll(Operation):
def __init__(self, shift, axis=None):
super().__init__()
self.shift = shift
self.axis = axis
def call(self, x):
return backend.execute("roll", x, self.shift, self.axis)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def roll(x, shift, axis=None):
if any_symbolic_tensors((x,)):
return Roll(shift, axis=axis).symbolic_call(x)
return backend.execute("roll", x, shift, axis=axis)
class Round(Operation):
def __init__(self, decimal=0):
super().__init__()
self.decimal = decimal
def call(self, x):
return backend.execute("round", x, self.decimal)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
class Sign(Operation):
def call(self, x):
return backend.execute("sign", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype="int32")
def sign(x):
if any_symbolic_tensors((x,)):
return Sign().symbolic_call(x)
return backend.execute("sign", x)
class Sin(Operation):
def call(self, x):
return backend.execute("sin", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape)
def sin(x):
if any_symbolic_tensors((x,)):
return Sin().symbolic_call(x)
return backend.execute("sin", x)
class Size(Operation):
def call(self, x):
return backend.execute("size", x)
def compute_output_spec(self, x):
return KerasTensor([], dtype="int32")
def size(x):
if any_symbolic_tensors((x,)):
return Size().symbolic_call(x)
return backend.execute("size", x)
class Sort(Operation):
def __init__(self, axis=-1):
super().__init__()
self.axis = axis
def call(self, x):
return backend.execute("sort", x, axis=self.axis)
def compute_output_spec(self, x):
return KerasTensor(x.shape, x.dtype)
def sort(x, axis=-1):
if any_symbolic_tensors((x,)):
return Sort(axis=axis).symbolic_call(x)
return backend.execute("sort", x, axis=axis)
class Split(Operation):
def __init__(self, indices_or_sections, axis=0):
super().__init__()
self.indices_or_sections = indices_or_sections
self.axis = axis
def call(self, x):
return backend.execute(
"split", x, self.indices_or_sections, axis=self.axis
)
def compute_output_spec(self, x):
x_shape = list(x.shape)
x_size_on_axis = x_shape[self.axis]
if isinstance(self.indices_or_sections, int):
if x_size_on_axis is None:
x_shape[self.axis] = None
return [
KerasTensor(x_shape, dtype=x.dtype)
for _ in range(self.indices_or_sections)
]
if np.mod(x_size_on_axis, self.indices_or_sections) != 0:
raise ValueError(
"`x` size on given `axis` must be dividible by "
"`indices_or_sections` when `indices_or_sections` is an "
f"int. But received {x_size_on_axis} and "
f"{self.indices_or_sections}."
)
size = x_size_on_axis // self.indices_or_sections
x_shape[self.axis] = size
return [
KerasTensor(x_shape, dtype=x.dtype)
for _ in range(self.indices_or_sections)
]
indices_or_sections = [0] + self.indices_or_sections
output_size = np.diff(indices_or_sections)
outputs = []
for i in range(len(output_size)):
output_shape = list(x_shape)
output_shape[self.axis] = int(output_size[i])
outputs.append(KerasTensor(output_shape, dtype=x.dtype))
return outputs
def split(x, indices_or_sections, axis=0):
if any_symbolic_tensors((x,)):
return Split(indices_or_sections, axis=axis).symbolic_call(x)
return backend.execute("split", x, indices_or_sections, axis=axis)
class Stack(Operation):
def __init__(self, axis=0):
super().__init__()
self.axis = axis
def call(self, xs):
return backend.execute("stack", xs, axis=self.axis)
def compute_output_spec(self, xs):
first_shape = xs[0].shape
for x in xs:
if not shape_equal(x.shape, first_shape, axis=[], allow_none=True):
raise ValueError(
"Every value in `xs` must have the same shape. But found "
f"element of shape {x.shape}, which is different from the "
f"first element's shape {first_shape}."
)
size_on_axis = len(xs)
output_shape = list(first_shape)
if self.axis == -1:
output_shape = output_shape + [size_on_axis]
elif self.axis >= 0:
output_shape.insert(self.axis, size_on_axis)
else:
output_shape.insert(self.axis + 1, size_on_axis)
return KerasTensor(output_shape, dtype=x.dtype)
def stack(x, axis=0):
if any_symbolic_tensors((x,)):
return Stack(axis=axis).symbolic_call(x)
return backend.execute("stack", x, axis=axis)
class Std(Operation):
def __init__(self, axis=None, keepdims=False):
super().__init__()
if isinstance(axis, int):
self.axis = [axis]
else:
self.axis = axis
self.keepdims = keepdims
def call(self, x):
return backend.execute("std", x, axis=self.axis, keepdims=self.keepdims)
def compute_output_spec(self, x):
return KerasTensor(
reduce_shape(x.shape, axis=self.axis, keepdims=self.keepdims),
)
def std(x, axis=None, keepdims=False):
if any_symbolic_tensors((x,)):
return Std(axis=axis, keepdims=keepdims).symbolic_call(x)
return backend.execute("std", x, axis=axis, keepdims=keepdims)
class Swapaxes(Operation):
def __init__(self, axis1, axis2):
super().__init__()
self.axis1 = axis1
self.axis2 = axis2
def call(self, x):
return backend.execute("swapaxes", x, self.axis1, self.axis2)
def compute_output_spec(self, x):
x_shape = list(x.shape)
tmp = x_shape[self.axis1]
x_shape[self.axis1] = x_shape[self.axis2]
x_shape[self.axis2] = tmp
return KerasTensor(x_shape, dtype=x.dtype)
def swapaxes(x, axis1, axis2):
if any_symbolic_tensors((x,)):
return Swapaxes(axis1, axis2).symbolic_call(x)
return backend.execute("swapaxes", x, axis1=axis1, axis2=axis2)
class Take(Operation):
def __init__(self, axis=None):
super().__init__()
self.axis = axis
def call(self, x, indices):
return backend.execute("take", x, indices, axis=self.axis)
def compute_output_spec(self, x, indices):
x_shape = list(x.shape)
indices_shape = list(getattr(np.array(indices), "shape", []))
if self.axis is None:
return KerasTensor(indices_shape, dtype=x.dtype)
if self.axis == -1:
output_shape = x_shape[:-1] + indices_shape
else:
output_shape = (
x_shape[: self.axis] + indices_shape + x_shape[self.axis + 1 :]
)
return KerasTensor(output_shape, dtype=x.dtype)
def take(x, indices, axis=None):
if any_symbolic_tensors((x, indices)):
return Take(axis=axis).symbolic_call(x, indices)
return backend.execute("take", x, indices, axis=axis)
class TakeAlongAxis(Operation):
def __init__(self, axis=None):
super().__init__()
self.axis = axis
def call(self, x, indices):
return backend.execute("take_along_axis", x, indices, axis=self.axis)
def compute_output_spec(self, x, indices):
x_shape = list(x.shape)
indices_shape = list(indices.shape)
if self.axis is None:
x_shape = [None] if None in x_shape else [int(np.prod(x_shape))]
if len(x_shape) != len(indices_shape):
raise ValueError(
"`x` and `indices` must have the same number of dimensions, "
f"but receive shape {x_shape} and {indices_shape}."
)
del x_shape[self.axis]
del indices_shape[self.axis]
output_shape = broadcast_shapes(x_shape, indices_shape)
size_on_axis = indices.shape[self.axis]
if self.axis == -1:
output_shape = output_shape + [size_on_axis]
elif self.axis >= 0:
output_shape.insert(self.axis, size_on_axis)
else:
output_shape.insert(self.axis + 1, size_on_axis)
return KerasTensor(output_shape, dtype=x.dtype)
def take_along_axis(x, indices, axis=None):
if any_symbolic_tensors((x, indices)):
return TakeAlongAxis(axis=axis).symbolic_call(x, indices)
return backend.execute("take_along_axis", x, indices, axis=axis)
class Tan(Operation):
def call(self, x):
return backend.execute("tan", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape)
def tan(x):
if any_symbolic_tensors((x,)):
return Tan().symbolic_call(x)
return backend.execute("tan", x)
class Tensordot(Operation):
def __init__(self, axes=2):
super().__init__()
self.axes = axes
def call(self, x1, x2):
return backend.execute("tensordot", x1, x2, axes=self.axes)
def compute_output_spec(self, x1, x2):
x1_shape = list(getattr(x1, "shape", []))
x2_shape = list(getattr(x2, "shape", []))
if not isinstance(self.axes, int):
x1_select_shape = [x1_shape[ax] for ax in self.axes[0]]
x2_select_shape = [x2_shape[ax] for ax in self.axes[1]]
if not shape_equal(
x1_select_shape, x2_select_shape, allow_none=True
):
raise ValueError(
"Shape mismatch on `x1[axes[0]]` and `x2[axes[1]]`, "
f"received {x1_select_shape} and {x2_select_shape}."
)
for ax in self.axes[0]:
x1_shape[ax] = -1
for ax in self.axes[1]:
x2_shape[ax] = -1
x1_shape = list(filter((-1).__ne__, x1_shape))
x2_shape = list(filter((-1).__ne__, x2_shape))
output_shape = x1_shape + x2_shape
return KerasTensor(output_shape, dtype=x1.dtype)
if self.axes <= 0:
output_shape = x1_shape + x2_shape
else:
output_shape = x1_shape[: -self.axes] + x2_shape[self.axes :]
return KerasTensor(output_shape, dtype=x1.dtype)
def tensordot(x1, x2, axes=2):
if any_symbolic_tensors((x1, x2)):
return Tensordot(axes=axes).symbolic_call(x1, x2)
return backend.execute("tensordot", x1, x2, axes=axes)
def round(x):
if any_symbolic_tensors((x,)):
return Round().symbolic_call(x)
return backend.execute("round", x)
class Tile(Operation):
def __init__(self, repeats):
super().__init__()
self.repeats = repeats
def call(self, x):
return backend.execute("tile", x, self.repeats)
def compute_output_spec(self, x):
x_shape = list(x.shape)
repeats = self.repeats
if len(x_shape) > len(repeats):
repeats = [1] * (len(x_shape) - len(repeats)) + repeats
else:
x_shape = [1] * (len(repeats) - len(x_shape)) + x_shape
output_shape = []
for x_size, repeat in zip(x_shape, repeats):
if x_size is None:
output_shape.append(None)
else:
output_shape.append(x_size * repeat)
return KerasTensor(output_shape, dtype=x.dtype)
def tile(x, repeats):
if any_symbolic_tensors((x,)):
return Tile(
repeats,
).symbolic_call(x)
return backend.execute("tile", x, repeats)
class Trace(Operation):
def __init__(self, offset=0, axis1=0, axis2=1):
super().__init__()
self.offset = offset
self.axis1 = axis1
self.axis2 = axis2
def call(self, x):
return backend.execute(
"trace", x, offset=self.offset, axis1=self.axis1, axis2=self.axis2
)
def compute_output_spec(self, x):
x_shape = list(x.shape)
x_shape[self.axis1] = -1
x_shape[self.axis2] = -1
output_shape = list(filter((-1).__ne__, x_shape))
return KerasTensor(output_shape, dtype=x.dtype)
def trace(x, offset=0, axis1=0, axis2=1):
if any_symbolic_tensors((x,)):
return Trace(offset, axis1, axis2).symbolic_call(x)
return backend.execute("trace", x, offset=offset, axis1=axis1, axis2=axis2)
class Tri(Operation):
def call(self, N, M=None, k=0, dtype="float32"):
return backend.execute("tri", N, M=M, k=k, dtype=dtype)
def compute_output_spec(self, N, M=None, k=0, dtype="float32"):
if M is None:
M = N
return KerasTensor((N, M), dtype=dtype)
def tri(N, M=None, k=0, dtype="float32"):
return backend.execute("tri", N, M=M, k=k, dtype=dtype)
class Tril(Operation):
def __init__(self, k=0):
super().__init__()
self.k = k
def call(self, x):
return backend.execute("tril", x, k=self.k)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def tril(x, k=0):
if any_symbolic_tensors((x,)):
return Tril(k=k).symbolic_call(x)
return backend.execute("tril", x, k=k)
class Triu(Operation):
def __init__(self, k=0):
super().__init__()
self.k = k
def call(self, x):
return backend.execute("triu", x, k=self.k)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def triu(x, k=0):
if any_symbolic_tensors((x,)):
return Triu(k=k).symbolic_call(x)
return backend.execute("triu", x, k=k)
class Vdot(Operation):
def call(self, x1, x2):
return backend.execute("vdot", x1, x2)
def compute_output_spec(self, x1, x2):
return KerasTensor([], dtype=x1.dtype)
def vdot(x1, x2):
if any_symbolic_tensors((x1, x2)):
return Vdot().symbolic_call(x1, x2)
return backend.execute("vdot", x1, x2)
class Vstack(Operation):
def call(self, xs):
return backend.execute("vstack", xs)
def compute_output_spec(self, xs):
first_shape = xs[0].shape
total_size_on_axis = 0
for x in xs:
if not shape_equal(x.shape, first_shape, axis=[0], allow_none=True):
raise ValueError(
"Every value in `xs` must have the same shape except on "
f"the `axis` dim. But found element of shape {x.shape}, "
f"which is different from the first element's "
f"shape {first_shape}."
)
if total_size_on_axis is None or x.shape[0] is None:
total_size_on_axis = None
else:
total_size_on_axis += x.shape[0]
output_shape = list(first_shape)
output_shape[0] = total_size_on_axis
return KerasTensor(output_shape)
def vstack(xs):
if any_symbolic_tensors((xs,)):
return Vstack().symbolic_call(xs)
return backend.execute("vstack", xs)
class Where(Operation):
def call(self, condition, x1, x2):
return backend.execute("where", condition, x1, x2)
def compute_output_spec(self, condition, x1, x2):
condition_shape = getattr(condition, "shape", [])
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
output_shape = broadcast_shapes(condition_shape, x1_shape)
output_shape = broadcast_shapes(output_shape, x2_shape)
return KerasTensor(output_shape, dtype=x1.dtype)
def where(condition, x1, x2):
if any_symbolic_tensors((condition, x1, x2)):
return Where().symbolic_call(condition, x1, x2)
return backend.execute("where", condition, x1, x2)
class Subtract(Operation):
def call(self, x1, x2):
return backend.numpy.subtract(x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
output_shape = broadcast_shapes(x1_shape, x2_shape)
return KerasTensor(output_shape, dtype=x1.dtype)
def subtract(x1, x2):
if any_symbolic_tensors((x1, x2)):
return Subtract().symbolic_call(x1, x2)
return backend.numpy.subtract(x1, x2)
class Multiply(Operation):
def call(self, x1, x2):
return backend.numpy.multiply(x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
output_shape = broadcast_shapes(x1_shape, x2_shape)
return KerasTensor(output_shape, dtype=x1.dtype)
def multiply(x1, x2):
if any_symbolic_tensors((x1, x2)):
return Multiply().symbolic_call(x1, x2)
return backend.numpy.multiply(x1, x2)
class Divide(Operation):
def call(self, x1, x2):
return backend.execute("divide", x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
output_shape = broadcast_shapes(x1_shape, x2_shape)
return KerasTensor(output_shape, dtype=x1.dtype)
def divide(x1, x2):
if any_symbolic_tensors((x1, x2)):
return Divide().symbolic_call(x1, x2)
return backend.execute("divide", x1, x2)
class TrueDivide(Operation):
def call(self, x1, x2):
return backend.execute("true_divide", x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
output_shape = broadcast_shapes(x1_shape, x2_shape)
return KerasTensor(output_shape, dtype=x1.dtype)
def true_divide(x1, x2):
if any_symbolic_tensors((x1, x2)):
return TrueDivide().symbolic_call(x1, x2)
return backend.execute("true_divide", x1, x2)
class Power(Operation):
def call(self, x1, x2):
return backend.execute("power", x1, x2)
def compute_output_spec(self, x1, x2):
x1_shape = getattr(x1, "shape", [])
x2_shape = getattr(x2, "shape", [])
output_shape = broadcast_shapes(x1_shape, x2_shape)
return KerasTensor(output_shape, dtype=x1.dtype)
def power(x1, x2):
if any_symbolic_tensors((x1, x2)):
return Power().symbolic_call(x1, x2)
return backend.execute("power", x1, x2)
class Negative(Operation):
def call(self, x):
return backend.execute("negative", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def negative(x):
if any_symbolic_tensors((x,)):
return Negative().symbolic_call(x)
return backend.execute("negative", x)
class Square(Operation):
def call(self, x):
return backend.execute("square", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def square(x):
if any_symbolic_tensors((x,)):
return Square().symbolic_call(x)
return backend.execute("square", x)
class Sqrt(Operation):
def call(self, x):
return backend.execute("sqrt", x)
def compute_output_spec(self, x):
return KerasTensor(x.shape, dtype=x.dtype)
def sqrt(x):
if any_symbolic_tensors((x,)):
return Sqrt().symbolic_call(x)
return backend.execute("sqrt", x)
class Squeeze(Operation):
def __init__(self, axis=None):
super().__init__()
self.axis = axis
def call(self, x):
return backend.execute("squeeze", x, axis=self.axis)
def compute_output_spec(self, x, axis=None):
input_shape = list(x.shape)
if axis is None:
output_shape = list(filter((1).__ne__, input_shape))
return KerasTensor(output_shape)
else:
if input_shape[axis] != 1:
raise ValueError(
f"Cannot squeeze axis {axis}, because the dimension is not "
"1."
)
del input_shape[axis]
return KerasTensor(input_shape, dtype=x.dtype)
def squeeze(x, axis=None):
if any_symbolic_tensors((x,)):
return Squeeze().symbolic_call(x, axis=axis)
return backend.execute("squeeze", x, axis=axis)
class Transpose(Operation):
def __init__(self, axes=None):
super().__init__()
self.axes = axes
def call(self, x):
return backend.execute("transpose", x, axes=self.axes)
def compute_output_spec(self, x):
x_shape = x.shape
if self.axes is None:
return KerasTensor(x_shape[::-1])
if len(self.axes) != len(x_shape):
raise ValueError(
"axis must be a list of the same length as the input shape, "
f"expected {len(x_shape)}, but received {len(self.axes)}."
)
output_shape = []
for ax in self.axes:
output_shape.append(x_shape[ax])
return KerasTensor(output_shape, dtype=x.dtype)
def transpose(x, axes=None):
if any_symbolic_tensors((x,)):
return Transpose(axes=axes).symbolic_call(x)
return backend.execute("transpose", x, axes=axes)
class Mean(Operation):
def __init__(self, axis=None, keepdims=False):
super().__init__()
if isinstance(axis, int):
axis = [axis]
self.axis = axis
self.keepdims = keepdims
def call(self, x):
return backend.numpy.mean(x, axis=self.axis, keepdims=self.keepdims)
def compute_output_spec(self, x):
return KerasTensor(
reduce_shape(x.shape, axis=self.axis, keepdims=self.keepdims),
dtype=x.dtype,
)
def mean(x, axis=None, keepdims=False):
if any_symbolic_tensors((x,)):
return Mean(axis=axis, keepdims=keepdims).symbolic_call(x)
return backend.numpy.mean(x, axis=axis, keepdims=keepdims)
class Var(Operation):
def __init__(self, axis=None, keepdims=False):
super().__init__()
if isinstance(axis, int):
axis = [axis]
self.axis = axis
self.keepdims = keepdims
def call(self, x):
return backend.execute("var", x, axis=self.axis, keepdims=self.keepdims)
def compute_output_spec(self, x):
return KerasTensor(
reduce_shape(x.shape, axis=self.axis, keepdims=self.keepdims),
dtype=x.dtype,
)
def var(x, axis=None, keepdims=False):
if any_symbolic_tensors((x,)):
return Var(axis=axis, keepdims=keepdims).symbolic_call(x)
return backend.execute("var", x, axis=axis, keepdims=keepdims)
class Sum(Operation):
def __init__(self, axis=None, keepdims=False):
super().__init__()
if isinstance(axis, int):
axis = [axis]
self.axis = axis
self.keepdims = keepdims
def call(self, x):
return backend.execute("sum", x, axis=self.axis, keepdims=self.keepdims)
def compute_output_spec(self, x):
return KerasTensor(
reduce_shape(x.shape, axis=self.axis, keepdims=self.keepdims),
dtype=x.dtype,
)
def sum(x, axis=None, keepdims=False):
if any_symbolic_tensors((x,)):
return Sum(axis=axis, keepdims=keepdims).symbolic_call(x)
return backend.execute("sum", x, axis=axis, keepdims=keepdims)
class Zeros(Operation):
def call(self, shape, dtype="float32"):
return backend.execute("zeros", shape, dtype)
def compute_output_spec(self, shape, dtype="float32"):
return KerasTensor(shape, dtype=dtype)
def zeros(shape, dtype="float32"):
return backend.execute("zeros", shape, dtype)
class Ones(Operation):
def call(self, shape, dtype="float32"):
return backend.execute("ones", shape, dtype)
def compute_output_spec(self, shape, dtype="float32"):
return KerasTensor(shape, dtype=dtype)
def ones(shape, dtype="float32"):
return backend.execute("ones", shape, dtype)
class Eye(Operation):
def call(self, N, M=None, k=0, dtype="float32"):
return backend.execute("eye", N, M=M, k=k, dtype=dtype)
def compute_output_spec(self, N, M=None, k=0, dtype="float32"):
if M is None:
M = N
return KerasTensor((N, M), dtype=dtype)
def eye(N, M=None, k=0, dtype="float32"):
return backend.execute("eye", N, M=M, k=k, dtype=dtype)
Reference in New Issue View Git Blame Copy Permalink