add erf and erfc rules

jax/experimental/jet.py

@@ -211,7 +211,36 @@ def linear_prop(prim, primals_in, series_in, **params):
 deflinear(lax_fft.fft_p)
 deflinear(xla.device_put_p)
 
+def def_deriv(prim, deriv):
+  """
+  Define the jet rule for a primitive in terms of its first derivative.
+  """
+  jet_rules[prim] = partial(deriv_prop, prim, deriv)
 
+def deriv_prop(prim, deriv, primals_in, series_in):
+  x, = primals_in
+  series, = series_in
+  primal_out = prim.bind(x)
+  c0, cs = jet(deriv, primals_in, series_in)
+  c = [c0] + cs
+  u = [x] + series
+  v = [primal_out] + [None] * len(series)
+  for k in range(1, len(v)):
+    v[k] = fact(k-1) * sum(_scale(k, j) * c[k-j] * u[j] for j in range(1, k + 1))
+  primal_out, *series_out = v
+  return primal_out, series_out
+
+
+def_deriv(lax.erf_p, lambda x: lax.mul(lax._const(x, 2. / onp.sqrt(onp.pi)), lax.exp(lax.neg(lax.square(x)))))
+
+def def_comp(prim, comp):
+  """
+  Define the jet rule for a primitive in terms of a composition of simpler primitives.
+  """
+  jet_rules[prim] = partial(jet, comp)
+
+
+def_comp(lax.erfc_p, lambda x: 1 - lax.erf(x))
 
 ### More complicated rules
 

tests/jet_test.py

@@ -225,6 +225,10 @@ def test_asinh(self):      self.unary_check(lax.asinh, lims=[-100, 100])
   def test_acosh(self):      self.unary_check(lax.acosh, lims=[-100, 100])
   @jtu.skip_on_devices("tpu")
   def test_atanh(self):      self.unary_check(lax.atanh, lims=[-1, 1])
+  @jtu.skip_on_devices("tpu")
+  def test_erf(self):        self.unary_check(lax.erf)
+  @jtu.skip_on_devices("tpu")
+  def test_erfc(self):       self.unary_check(lax.erfc)
 
   @jtu.skip_on_devices("tpu")
   def test_div(self):   self.binary_check(lambda x, y: x / y, lims=[0.8, 4.0])