WEBVTT

1
00:00:00.000 --> 00:00:02.455
Part 8 of ML Libraries with a Bot.

2
00:00:02.455 --> 00:00:05.272
JAX on the CPU: Field Notes with a Bot.

3
00:00:05.272 --> 00:00:11.266
JAX takes NumPy-style code and adds two
tricks: automatic gradients and
compilation

4
00:00:11.266 --> 00:00:11.916
with jit.

5
00:00:11.916 --> 00:00:17.909
It's popular in research, and it's a good
library for learning how gradients really

6
00:00:17.909 --> 00:00:22.170
work, because you call grad on an ordinary
Python function.

7
00:00:22.170 --> 00:00:28.236
These field notes are what we saw
installing and running JAX on a CPU-only
box, with

8
00:00:28.236 --> 00:00:30.475
a bot explaining along the way.

9
00:00:31.292 --> 00:00:33.760
What the bot can and cannot do.

10
00:00:33.760 --> 00:00:36.546
Explain what grad, jit and vmap do.

11
00:00:36.546 --> 00:00:40.049
Convert a small NumPy function to
jax.numpy.

12
00:00:40.049 --> 00:00:44.508
Work through the calculus so you can check
JAX's answer.

13
00:00:44.508 --> 00:00:47.215
Keep up with JAX's install extras.

14
00:00:47.215 --> 00:00:49.444
Check the installation page.

15
00:00:49.444 --> 00:00:55.972
Know about JAX's sharp edges unless you
ask, such as immutable arrays and explicit

16
00:00:55.972 --> 00:00:56.927
random keys.

17
00:00:56.927 --> 00:00:59.077
Promise speed-ups from jit.

18
00:00:59.077 --> 00:01:00.112
Measure them.

19
00:01:00.792 --> 00:01:06.402
Run a function normally, check its answer
by hand, then add jit.

20
00:01:06.402 --> 00:01:12.275
If the answers differ, you've learned
something before it cost you.

21
00:01:12.875 --> 00:01:20.596
JAX's installation page says it needs two
packages: jax, which is pure Python, and

22
00:01:20.596 --> 00:01:24.268
jaxlib, which holds the compiled parts.

23
00:01:24.268 --> 00:01:30.764
For CPU, it gives: On our box that
installed jax and jaxlib 0.12.0.1.

24
00:01:31.458 --> 00:01:32.692
Real output:

25
00:01:33.292 --> 00:01:42.703
The loss is the mean of (w·x − y)² with x
= 1, 2, 3, 4, 5 and y = 3x.

26
00:01:42.703 --> 00:01:51.160
Its derivative with respect to w is the
mean of 2·x·(w·x − y).

27
00:01:51.160 --> 00:01:57.025
At w = 1, each term is 2·x·(x − 3x) =
−4x².

28
00:01:57.625 --> 00:02:05.635
jax.jit(loss) compiled the function, and
at w = 3 the loss is 0, because 3x matches

29
00:02:05.635 --> 00:02:06.600
y exactly.

30
00:02:06.600 --> 00:02:11.039
Compiled and uncompiled versions should
agree.

31
00:02:11.039 --> 00:02:16.637
We didn't time them here, so we make no
claim about speed.

32
00:02:17.333 --> 00:02:21.018
JAX doesn't use a hidden global random
state.

33
00:02:21.018 --> 00:02:25.357
You make a key with jax.random.key(0) and
pass it in.

34
00:02:25.357 --> 00:02:31.743
The same key gives the same numbers, so we
saw 1.623, 2.025, -0.434 (rounded).

35
00:02:31.743 --> 00:02:34.773
To get new numbers you split the key.

36
00:02:34.773 --> 00:02:37.966
JAX's Quickstart and docs explain this.

37
00:02:37.966 --> 00:02:41.650
Bots trained on NumPy habits often forget
it.

38
00:02:42.250 --> 00:02:45.156
In NumPy you can write a[0] = 5.

39
00:02:45.156 --> 00:02:51.514
In JAX that raises an error, and you write
a = a.at[0].set(5) instead.

40
00:02:51.514 --> 00:02:56.692
The jax.numpy reference lists this and
other differences.

41
00:02:57.292 --> 00:03:01.918
The prompt: Convert this NumPy function to
jax.numpy.

42
00:03:01.918 --> 00:03:04.973
List every change you made and why.

43
00:03:04.973 --> 00:03:12.130
If anything depends on in-place updates or
global random state, rewrite it the JAX

44
00:03:12.130 --> 00:03:12.479
way.

45
00:03:13.167 --> 00:03:18.807
If a bot tells you JAX will be faster for
your job, ask it how you would measure

46
00:03:18.807 --> 00:03:22.051
that on your own machine, and then measure
it.

47
00:03:22.051 --> 00:03:27.903
A fair test runs each version a few times
on the same data, and skips the first jit

48
00:03:27.903 --> 00:03:30.301
call, which includes compile time.

49
00:03:31.042 --> 00:03:37.149
JAX suits work where you want gradients of
your own functions, or compiled numerical

50
00:03:37.149 --> 00:03:39.113
code that looks like NumPy.

51
00:03:39.113 --> 00:03:45.148
If you're following a course or a model
that uses PyTorch or TensorFlow, stick
with

52
00:03:45.148 --> 00:03:46.529
that library first.

53
00:03:46.529 --> 00:03:52.273
A bot can translate small examples between
them, but check every translation by

54
00:03:52.273 --> 00:03:55.182
running both versions on the same input.

55
00:03:56.000 --> 00:03:57.902
Tested vs documented.

56
00:03:57.902 --> 00:04:05.330
Tested by us: the CPU install, the
gradient, jit result and random numbers,
on JAX

57
00:04:05.330 --> 00:04:06.145
0.12.0.1.

58
00:04:06.145 --> 00:04:13.210
Documented, not tested by us: GPU installs
and driver requirements, from JAX's

59
00:04:13.210 --> 00:04:14.840
installation page.

60
00:04:14.840 --> 00:04:18.011
Not claimed: any speed-up from jit.

61
00:04:18.792 --> 00:04:26.549
The last part today reads Hugging Face
model cards and licenses before you
download

62
00:04:26.549 --> 00:04:27.390
anything.

63
00:04:28.083 --> 00:04:34.127
Before you run anything: check every bot
answer against the official docs, work in
a

64
00:04:34.127 --> 00:04:40.027
virtual environment, load model files only
from sources you trust, respect dataset

65
00:04:40.027 --> 00:04:45.279
and model licenses, and keep personal data
out of bots and training sets.

66
00:04:45.279 --> 00:04:47.798
Every source is linked on the page.