def test_parens(): ExampleBias = np.random.randn(BATCH_SIZE) Input = np.random.randn(BATCH_SIZE, HIDDEN_DIM) Weight = np.random.randn(HIDDEN_DIM, OUTPUT_DIM) compiled = op( "Output[batch, output] = " "(Input[batch, hidden] + ExampleBias[batch]) * Weight[hidden, output]" ) result = compiled(Input=Input, ExampleBias=ExampleBias, Weight=Weight) simpler_compiled = op( "Output[batch, output] = Input[batch, hidden] * Weight[hidden, output]" ) simpler_result = simpler_compiled(Input=Input + np.expand_dims(ExampleBias, 1), Weight=Weight) assert np.allclose(result, simpler_result) expected = np.matmul(Input + np.expand_dims(ExampleBias, 1), Weight) manual = np.zeros((BATCH_SIZE, OUTPUT_DIM)) for batch in range(BATCH_SIZE): for hidden in range(HIDDEN_DIM): for output in range(OUTPUT_DIM): manual[batch, output] += (Input[batch, hidden] + ExampleBias[batch]) * Weight[hidden, output] assert result.shape == expected.shape assert np.allclose(manual, expected) assert np.allclose(result, manual) assert np.allclose(result, expected)
def test_empty_sum(): x = np.random.randn(5, 7) compiled = op("Output[] = x[a, b]") result = compiled(x=x) expected = np.sum(x, keepdims=False) assert result.shape == expected.shape assert np.allclose(result, expected)
def test_simple_parens(): ExampleBias = np.random.randn(BATCH_SIZE) Input = np.random.randn(BATCH_SIZE, HIDDEN_DIM) compiled = op("Output[batch, hidden] = (Input[batch, hidden] + ExampleBias[batch])") result = compiled(Input=Input, ExampleBias=ExampleBias) expected = Input + np.expand_dims(ExampleBias, 1) assert result.shape == expected.shape assert np.allclose(result, expected)
def test_broadcast_add(): A = np.random.randn(2) B = np.random.randn(3) expected = np.expand_dims(A, 1) + np.expand_dims(B, 0) broadcast_add = op("Y[a, b] = A[a] + B[b]") result = broadcast_add(A=A, B=B) assert result.shape == expected.shape assert np.allclose(result, expected)
def test_proper_bcast_reduce(): Input = np.random.randn(BATCH_SIZE, HIDDEN_DIM) ExampleBias = np.random.randn(BATCH_SIZE) OutputBias = np.random.randn(HIDDEN_DIM) compiled = op("Output[batch, hidden] = Input[batch, hidden] + ExampleBias[batch] + OutputBias[hidden]") result = compiled(Input=Input, ExampleBias=ExampleBias, OutputBias=OutputBias) expected = Input + np.expand_dims(ExampleBias, 1) + np.expand_dims(OutputBias, 0) assert result.shape == expected.shape assert np.allclose(result, expected)
def test_negation_subtraction(): X = np.random.randn(BATCH_SIZE, HIDDEN_DIM) W = np.random.randn(HIDDEN_DIM, OUTPUT_DIM) B = np.random.randn(OUTPUT_DIM) wx_plus_b = op("Y[i, k] = -W[j, k] * X[i, j] - B[k]") result = wx_plus_b(W=W, X=X, B=B) expected = np.matmul(-X, W) - B assert result.shape == expected.shape assert np.allclose(result, expected)
def test_y_equals_wx_plus_b(): X = np.random.randn(BATCH_SIZE, HIDDEN_DIM) W = np.random.randn(HIDDEN_DIM, OUTPUT_DIM) B = np.random.randn(OUTPUT_DIM) wx_plus_b = op("Y[i, k] = W[j, k] * X[i, j] + B[k]") result = wx_plus_b(W=W, X=X, B=B) expected = np.matmul(X, W) + B assert result.shape == expected.shape assert np.allclose(result, expected)
def test_matrix_decomp(): SMALLER_DIM = 2 X = np.random.randn(BATCH_SIZE, HIDDEN_DIM) W1 = np.random.randn(HIDDEN_DIM, SMALLER_DIM) W2 = np.random.randn(SMALLER_DIM, HIDDEN_DIM) compiled = op( "Output[batch, output] = X[batch, hidden] * W1[hidden, smaller] * W2[smaller, output]" ) result = compiled(X=X, W1=W1, W2=W2) expected = np.matmul(np.matmul(X, W1), W2) assert result.shape == expected.shape assert np.allclose(result, expected)
def test_more_complex(): ExampleBias = np.random.randn(BATCH_SIZE) Input = np.random.randn(BATCH_SIZE, HIDDEN_DIM) Weight = np.random.randn(HIDDEN_DIM, OUTPUT_DIM) OutputBias = np.random.randn(OUTPUT_DIM) compiled = op( "Output[batch, output] = " "(Input[batch, hidden] + ExampleBias[batch]) * Weight[hidden, output] + OutputBias[output]" ) result = compiled(Input=Input, ExampleBias=ExampleBias, Weight=Weight, OutputBias=OutputBias) expected = np.matmul(Input + np.expand_dims(ExampleBias, 1), Weight) + np.expand_dims(OutputBias, 0) assert result.shape == expected.shape assert np.allclose(result, expected)
def test_missing_tensor(): with pytest.raises(MissingTensor): X = np.random.randn(BATCH_SIZE, HIDDEN_DIM) matmul = op("Y[i, k] = X[i, j] * W[j, k]") result = matmul(X=X)
def test_unexpected_tensor(): with pytest.raises(UnexpectedTensor): X = np.random.randn(BATCH_SIZE, HIDDEN_DIM) W = np.random.randn(OUTPUT_DIM, OUTPUT_DIM) matmul = op("Y[i, k] = X[i, j] * W[j, k]") result = matmul(A=X)
def test_axis_disagreement(): with pytest.raises(AxisDisagreement): X = np.random.randn(BATCH_SIZE, HIDDEN_DIM) W = np.random.randn(OUTPUT_DIM, OUTPUT_DIM) matmul = op("Y[i, k] = X[i, j] * W[j, k]") result = matmul(X=X, W=W)
def test_matmul_transform(): X = np.random.randn(BATCH_SIZE, HIDDEN_DIM) W = np.random.randn(HIDDEN_DIM, OUTPUT_DIM) matmul = op("Y[i, k] = X[i, j] * W[j, k]") result = matmul(X=X, W=W) assert np.allclose(result, np.matmul(X, W))
def test_negation(): X = np.random.randn(HIDDEN_DIM) negate = op("Y[n] = -X[n]") result = negate(X=X) assert np.allclose(result, -X)