Python with_sharding_constraint Beispiele

Beispiel #1

 def f(x):
   y = jnp.dot(x, x)
   y = with_sharding_constraint(y, P(2,1))
   return y * 2

Beispiel #2

 def f(x):
   y = x + 1
   p, vjp_f = vjp(lambda z: jnp.sin(with_sharding_constraint(z, P(2,2))), y)
   return vjp_f(p)

Beispiel #3

 def f(x):
   y = x + 1
   y = with_sharding_constraint(y, P(2,1))
   return y * 2

Beispiel #4

 def f(x):
   return lax.while_loop(lambda i: i[0,0] < 10.,
                         lambda i: with_sharding_constraint(i + 1., P(2, 1)),
                         x)