Alleo
Alleo

Reputation: 8498

Defining expression in theano

I have a question about theano library. Suppose I have two matrices: a_ij and b_ij that respectively have shapes [n, m1] and [n, m2].

I want to define an expression which gives tensor c of shape [n, m1, m2]:

c_ijk = a_ij * b_ik

(pay attention that there is no sum, so this is not a dot product). Of course I can do this calculation in a loop, but I want theano to calculate derivatives for me.

Thanks in advance!

Upvotes: 1

Views: 53

Answers (1)

eickenberg
eickenberg

Reputation: 14377

Theano tensors broadcast exactly like numpy arrays. All you have to do is

import theano
import theano.tensor as T

import numpy as np
rng = np.random.RandomState(42)

n, m1, m2  = 10, 20, 30
A = T.fmatrix()
B = T.fmatrix()

C = A.reshape((n, m1, 1)) * B.reshape((n, 1, m2))

func_C = theano.function([A, B], C)

a = rng.randn(n, m1).astype(np.float32)
b = rng.randn(n, m2).astype(np.float32)

r1 = func_C(a, b)
r2 = np.einsum("ij, ik -> ijk", a, b)
assert np.sqrt(((r1 - r2) ** 2).sum()) < 1e-5

Where the definition of C is the crucial part. This can also be done with dynamic shapes, but this answer stays within your framework.

Upvotes: 1

Related Questions