Ejemplos de zeros_like_jaxval en Python

Ejemplo n.º 1

Archivo: ad.py Proyecto: jbampton/jax

 def join(self, xt, yt):
   xz, yz = type(xt) is Zero, type(yt) is Zero
   if xz == yz:
     return xt, yt
   elif yz and not xz:
     return xt, zeros_like_jaxval(xt)
   elif xz and not yz:
     return zeros_like_jaxval(yt), yt
   else:
     raise TypeError((xt, yt))

Ejemplo n.º 2

Archivo: batching.py Proyecto: xueeinstein/jax

def zeros_like_batched(batched_args, batch_dims):
  val, = batched_args
  bdim, = batch_dims
  return zeros_like_jaxval(val), bdim

Ejemplo n.º 3

Archivo: ad.py Proyecto: jbampton/jax

def replace_float0s(primal, tangent):
  if dtype(tangent) is float0:
    return zeros_like_jaxval(primal)
  else:
    return tangent