def test_unary(symbol, dims):
sizes = {'a': 3, 'b': 4}
shape = tuple(sizes[d] for d in dims)
inputs = OrderedDict((d, Bint[sizes[d]]) for d in dims)
dtype = 'real'
data = rand(shape) + 0.5
if symbol == '~':
data = ops.astype(data, 'uint8')
dtype = 2
if symbol == 'atanh':
data = ops.clamp(data, -0.99, 0.99)
if get_backend() != "torch" and symbol in [
"abs", "atanh", "sqrt", "exp", "log", "log1p", "sigmoid", "tanh"
]:
expected_data = getattr(ops, symbol)(data)
else:
expected_data = unary_eval(symbol, data)
x = Tensor(data, inputs, dtype)
actual = unary_eval(symbol, x)
check_funsor(actual, inputs, Array[dtype, ()], expected_data)
def einsum(equation, *operands):
"""
Log-sum-exp implementation of einsum.
"""
if get_backend() != "jax":
# NB: rename symbols to support NumPy, which allow only symbols a-z.
symbols = sorted(set(equation) - set(',->'))
rename = dict(zip(symbols, 'abcdefghijklmnopqrstuvwxyz'))
equation = ''.join(rename.get(s, s) for s in equation)
inputs, output = equation.split('->')
if inputs == output:
return operands[0][...] # create a new object
inputs = inputs.split(',')
shifts = []
exp_operands = []
for dims, operand in zip(inputs, operands):
shift = operand
for i, dim in enumerate(dims):
if dim not in output:
shift = ops.amax(shift, i, keepdims=True)
# avoid nan due to -inf - -inf
shift = ops.clamp(shift, ops.finfo(shift).min, None)
exp_operands.append(ops.exp(operand - shift))
# permute shift to match output
shift = shift.reshape(
[size for size, dim in zip(operand.shape, dims) if dim in output])
if len(shift.shape) > 0:
shift = shift.reshape((1, ) * (len(output) - shift.ndim) +
shift.shape)
dims = [dim for dim in dims if dim in output]
dims = [dim for dim in output if dim not in dims] + dims
shift = ops.permute(shift, [dims.index(dim) for dim in output])
shifts.append(shift)
result = ops.log(ops.einsum(equation, *exp_operands))
return sum(shifts + [result])
def clamp_finite(self):
finfo = ops.finfo(self.data)
data = ops.clamp(self.data, finfo.min, finfo.max)
return Tensor(data, self.inputs, self.dtype)