def single_state_reshape_same_state(inp: dace.float64[9], target_shape: dace.int64[2]): reshaped = dace.define_local([3, 3], dace.float64) donnx.ONNXReshape(data=inp, shape=target_shape, reshaped=reshaped) Zl = dace.elementwise(lambda x: log(x + 1), reshaped) S = np.sum(Zl) return S
def middle_sqrt_with_intermediate(Y: dace.float32[3, 3]): intermediate = dace.define_local([3, 3], dace.float32) W = dace.define_local([3, 3], dace.float32) intermediate[:] = dace.elementwise(lambda x: sqrt(x), Y) inner_sdfg_with_intermediate(intermediate, W) Z = np.sum(W) return Z
def outer_sqrt_with_intermediate(Y: dace.float32[3, 3]): intermediate = dace.define_local([3, 3], dace.float32) W = dace.define_local([3, 3], dace.float32) intermediate[:] = dace.elementwise(lambda x: sqrt(x), Y) W[:] = middle_sqrt_no_sum(intermediate) Z = np.sum(W) return Z
def dace_func(X: dace.float64[4, 5], Y: dace.float64[4, 3], W: dace.float64[7, 4, 3]): Xt[:] = np.transpose(X) YW[:] = np.min(W, axis=0) * Y Z[:] = Xt @ YW Zl = dace.elementwise(lambda x: log(x + 1), Z) S = np.sum(Zl) return S
def add_reshape_grad_test_nested(inp: dace.float64[9], bias: dace.float64[3], target_shape: dace.int64[2], result: dace.float64): reshaped = dace.define_local([3, 3], dace.float64) added = inp + 1 donnx.ONNXReshape(data=added, shape=target_shape, reshaped=reshaped) Z = reshaped * bias Zl = dace.elementwise(lambda x: log(x + 1), Z) result[:] = np.sum(Zl)
def reduce_node_1_axis_and_none_axis(X: dace.float32[4, 5], Y: dace.float32[4, 3], W: dace.float32[7, 4, 3]): Xt[:] = np.transpose(X) YW[:] = np.sum(W, axis=0) * Y Z[:] = Xt @ YW Zl = dace.elementwise(lambda x: log(x + 1), Z) S = np.sum(Zl) return S
def prog(A: dace.float64[5, 3], B: dace.float64[5, 3]): dace.elementwise(lambda x: log(x), A, B)
def elementwise_cast(A: dace.float32[5, 3], B: dace.float64[5, 3]): dace.elementwise(lambda x: x, A, B)
def elementise_none(A: dace.float64[5, 3], B: dace.float64[5, 3]): B[:] = dace.elementwise(lambda x: log(x), A)
def prog(A: dace.float32[5, 3], B: dace.float64[5, 3]): dace.elementwise(lambda x: x, A, B)
def inner_sdfg_with_intermediate(Z: dace.float32[3, 3], W: dace.float32[3, 3]): intermediate = dace.define_local([3, 3], dace.float32) intermediate[:] = dace.elementwise(lambda x: sqrt(x), Z) W[:] = dace.elementwise(lambda x: log(x), intermediate)
def inner_sdfg(Z: dace.float32[3, 3], W: dace.float32[3, 3]): W[:] = dace.elementwise(lambda x: log(x), Z)
def prog(input, output): output[:] = dace.elementwise(lambda x: log(x), input)
def prog(X, Y): Y[:] = dace.elementwise(cast_lambda, X)
def prog(X, Y): Y[:] = dace.elementwise(tanh_lambda, X)
def prog(X, Y): Y[:] = dace.elementwise(lambda x: sqrt(x), X)