<jupyter_start><jupyter_text>JaxTon    💯 JAX exercises                                                             This is Set 4: Just-In-Time (JIT) Compilation (Solutions 31-40) of JaxTon: 💯 JAX exercises        You can find all the exercises and solutions on GitHub **Prerequisites*** The configuration of jax should be set as shown in the code snippet below in order to use TPUs.* A sample function `cube` will be used for the exercises.<jupyter_code>!python3 -m pip install jax
import jax
import jax.numpy as jnp
import os
import requests

try:
    url = 'http:' + os.environ['TPU_NAME'].split(':')[1] + ':8475/requestversion/tpu_driver_nightly'
    resp = requests.post(url)
    jax.config.FLAGS.jax_xla_backend = 'tpu_driver'
    jax.config.FLAGS.jax_backend_target = os.environ['TPU_NAME']
except:
    pass

jax.devices()
def cube(x):
    return x**3

cube(2.1)<jupyter_output><empty_output><jupyter_text>**Exercise 31: JIT-compile the `cube` function and assign it to `cube_jit`**<jupyter_code>cube_jit = jax.jit(cube)
cube_jit<jupyter_output><empty_output><jupyter_text>**Exercise 32: Display execution time of `cube_jit` for first run (with overhead) with input=10.24**<jupyter_code>%%time
cube_jit(10.24)<jupyter_output>CPU times: user 221 ms, sys: 149 ms, total: 370 ms
Wall time: 366 ms<jupyter_text>**Exercise 33: Display execution time of `cube_jit` for second run (without overhead) with input=10.24**<jupyter_code>%%time
cube_jit(10.24)<jupyter_output>CPU times: user 16.8 ms, sys: 20.5 ms, total: 37.2 ms
Wall time: 37.4 ms<jupyter_text>**Exercise 34: Run `cube_jit` with input=10.24 and assign it to `cube_value`**<jupyter_code>cube_value = cube_jit(10.24)
cube_value<jupyter_output><empty_output><jupyter_text>**Exercise 35: Run `cube_jit` with jit disabled and input=10.24 and assign it to `cube_value_nojit`**<jupyter_code>with jax.disable_jit():
    cube_value_nojit = cube_jit(10.24)

cube_value_nojit<jupyter_output><empty_output><jupyter_text>**Exercise 36: Evaluate the shape of `cube_jit` with input=10.24 and assign it to `cube_shape`**<jupyter_code>cube_shape = jax.eval_shape(cube_jit, 10.24)
cube_shape<jupyter_output><empty_output><jupyter_text>**Exercise 37: Create the jaxpr of `cube_jit` with input=10.24 and assign it to `cube_jaxpr`**<jupyter_code>cube_jaxpr = jax.make_jaxpr(cube_jit)(10.24)
cube_jaxpr<jupyter_output><empty_output><jupyter_text>**Exercise 38: Assign the XLA computation of `cube_jit` with input=10.24 to `cube_xla` and print it's XLA HLO text**<jupyter_code>cube_xla = jax.xla_computation(cube_jit)(10.24)
print(cube_xla.as_hlo_text())<jupyter_output>HloModule xla_computation_cube.12

jit_cube__1.3 {
  constant.5 = pred[] constant(false)
  parameter.4 = f32[] parameter(0)
  multiply.6 = f32[] multiply(parameter.4, parameter.4)
  multiply.7 = f32[] multiply(parameter.4, multiply.6)
  ROOT tuple.8 = (f32[]) tuple(multiply.7)
}

ENTRY xla_computation_cube.12 {
  constant.2 = pred[] constant(false)
  parameter.1 = f32[] parameter(0)
  call.9 = (f32[]) call(parameter.1), to_apply=jit_cube__1.3
  get-tuple-element.10 = f32[] get-tuple-element(call.9), index=0
  ROOT tuple.11 = (f32[]) tuple(get-tuple-element.10)
}<jupyter_text>**Exercise 39: Use the name `jaxton_cube_fn` internally for the `cube_jit` function and assign the named function to `cube_named_jit`**<jupyter_code>cube_named_jit = jax.named_call(cube_jit, name='jaxton_cube_fn')
cube_named_jit<jupyter_output><empty_output><jupyter_text>**Exercise 40: Assign the XLA computation of `cube_named_jit` with input=10.24 to `cube_named_xla` and print it's XLA HLO text**<jupyter_code>cube_named_xla = jax.xla_computation(cube_named_jit)(10.24)
print(cube_named_xla.as_hlo_text())<jupyter_output>HloModule xla_computation_cube__1.18

jit_cube__2.3 {
  constant.5 = pred[] constant(false)
  parameter.4 = f32[] parameter(0)
  multiply.6 = f32[] multiply(parameter.4, parameter.4)
  multiply.7 = f32[] multiply(parameter.4, multiply.6)
  ROOT tuple.8 = (f32[]) tuple(multiply.7)
}

jaxton_cube_fn.9 {
  constant.11 = pred[] constant(false)
  parameter.10 = f32[] parameter(0)
  call.12 = (f32[]) call(parameter.10), to_apply=jit_cube__2.3
  get-tuple-element.13 = f32[] get-tuple-element(call.12), index=0
  ROOT tuple.14 = (f32[]) tuple(get-tuple-element.13)
}

ENTRY xla_computation_cube__1.18 {
  constant.2 = pred[] constant(false)
  parameter.1 = f32[] parameter(0)
  call.15 = (f32[]) call(parameter.1), to_apply=jaxton_cube_fn.9
  get-tuple-element.16 = f32[] get-tuple-element(call.15), index=0
  ROOT tuple.17 = (f32[]) tuple(get-tuple-element.16)
}
{'fraction_non_alphanumeric': 0.1309289617486339, 'fraction_numerical': 0.0546448087431694, 'mean_word_length': 5.490780141843971, 'pattern_counts': {'":': 0, '<': 41, '<?xml version=': 0, '>': 41, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '190535', 'license': 'Apache-2.0', 'n_tokens_mistral': 1967, 'n_tokens_neox': 1819, 'n_words': 377, 'path': 'notebooks/04_justintime_jit_compilation_solutions.ipynb', 'repo_name': 'nirmalya8/jaxton', 'size': '4575'}