Implement a Jax conversion for the Second Op (#185)

tests/sandbox/test_jax.py

@@ -696,6 +696,20 @@ def test_identity():
     compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])
 
 
+def test_second():
+    a0 = tt.scalar("a0")
+    b = tt.scalar("b")
+
+    out = theano.scalar.basic.second(a0, b)
+    fgraph = theano.gof.FunctionGraph([a0, b], [out])
+    compare_jax_and_py(fgraph, [10.0, 5.0])
+
+    a1 = tt.vector("a1")
+    out = tt.second(a1, b)
+    fgraph = theano.gof.FunctionGraph([a1, b], [out])
+    compare_jax_and_py(fgraph, [np.zeros([5], dtype=theano.config.floatX), 5.0])
+
+
 def test_shared():
     a = theano.shared(np.array([1, 2, 3], dtype=theano.config.floatX))
 

theano/sandbox/jaxify.py

@@ -19,7 +19,7 @@ from theano.compile.ops import (
 )
 from theano.gof import FunctionGraph
 from theano.ifelse import IfElse
-from theano.scalar.basic import Cast, Clip, Composite, Identity, ScalarOp
+from theano.scalar.basic import Cast, Clip, Composite, Identity, ScalarOp, Second
 from theano.scan.op import Scan
 from theano.scan.utils import scan_args as ScanArgs
 from theano.tensor.basic import (
@@ -256,6 +256,14 @@ def jax_funcify_ScalarSoftplus(op):
     return scalarsoftplus
 
 
+@jax_funcify.register(Second)
+def jax_funcify_Second(op):
+    def second(x, y):
+        return jnp.broadcast_to(y, x.shape)
+
+    return second
+
+
 @jax_funcify.register(AllocEmpty)
 def jax_funcify_AllocEmpty(op):
     def allocempty(*shape):