Programming TPUs in JAX | How To Scale Your Model

[jacobaustin123]

JAX programming, broadly construed!

[eitanporat]

Hi! Thanks for the great book.

I believe that Problem 1's attached solution doesn't correspond to the problem description (it's only on one axis).

I have been getting an issue with my solution and have written a Github Issue. I would really appreciate if you could help me out

I also have gotten another error I couldn't solve. I run this line (as described in this chapter):

  replicated = jax.lax.with_sharding_constraint(replicated, P(None, 'Y'))

and I keep getting an error:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
/tmp/ipython-input-215-3887179796.py in <cell line: 0>()
     10 
     11 with mesh:
---> 12   average_A_shmap = average_shard(A)
     13 
     14 average_A_shmap.addressable_shards

    [... skipping hidden 13 frame]

1 frames
    [... skipping hidden 15 frame]

/usr/local/lib/python3.11/dist-packages/jax/_src/util.py in __hash__(self)
    520 
    521   def __hash__(self):
--> 522     return hash(
    523       (
    524         self.f.__code__,

TypeError: unhashable type: 'set'

has the API perhaps changed?

[eitanporat]

For more context, here is my Colab notebook.

[jacobaustin123]

Hi Eitan. The question wasn't super well phrased and had a few mistakes in it, so thank you for pointing this out. I've gone ahead and clarified the question and also given full solutions to both parts. Let me know if this makes more sense.

I think the sharding constraint may need to be a NamedSharding and not PartitionSpec. I've also updated the book.