!pip install --upgrade git+https://github.com/google/learned_optimization.git oryx tensorflow==2.8.0rc0 numpy

"""
Explanation: Summary tutorial: Getting metrics out of your models
The goal of the learned_optimization.summary module is to seamlessly allow researchers to annotate and extract data from within a jax computation / machine learning model.
This could be anything from mean and standard deviation of an activation, to looking at the distribution of outputs.
Doing this in Jax can be challenging at times as code written in Jax can make use of a large number of function transformations making it difficult to reach in and look at a value.
This notebook discusses the learned_optimization.summary module which provides one solution to this problem and is discussed in this notebook.
Deps
In addition to learned_optimization, the summary module requires oryx. This can be a bit finicky to install at the moment as it relies upon particular versions of tensorflow (even though we never use these pieces).
All of learned_optimization will run without oryx, but to get summaries this module must be installed.
In a colab this is even more annoying as we must first upgrade the versions of some installed modules, restart the colab kernel, and then proceed to run the remainder of the cells.
End of explanation
"""


import oryx
from learned_optimization import summary
assert summary.ORYX_LOGGING

"""
Explanation: To check that everything is operating as expected we can check that the imports succeed.
End of explanation
"""


import jax
import jax.numpy as jnp

def forward(params):
  to_look_at = jnp.mean(params) * 2.
  return params


def loss(parameters):
  loss = jnp.mean(forward(parameters)**2)
  return loss


value_grad_fn = jax.jit(jax.value_and_grad(loss))
value_grad_fn(1.0)

"""
Explanation: Basic Example
Let's say that we have the following function and we wish to look at the to_look_at value.
End of explanation
"""


def forward(params):
  to_look_at = jnp.mean(params) * 2.
  summary.summary("to_look_at", to_look_at)
  return params


@jax.jit
def loss(parameters):
  loss = jnp.mean(forward(parameters)**2)
  return loss

"""
Explanation: With the summary module we can first annotate this value with summary.summary
End of explanation
"""


result, metrics = summary.with_summary_output_reduced(loss)(1.)
result, metrics

"""
Explanation: Then we can transform the loss function with the function transformation: summary.with_summary_output_reduced.
This transformation goes through the computation and extracts all the tagged values and returns them to us by name in a dictionary.
In implementation, all the hard work here is done by the wonderful oryx library (in particular harvest).
When we wrap a function this, we return a tuple containing the original result, and a dictionary with the desired metrics.
End of explanation
"""


summary.aggregate_metric_list([metrics])

"""
Explanation: As currently returned, the dictionary contains extra information as well as potentially duplicate values. We can collapse these metrics into a single value with the following:
End of explanation
"""


result, metrics = jax.jit(summary.with_summary_output_reduced(loss))(1.)
summary.aggregate_metric_list([metrics])

"""
Explanation: The keys of this dictionary first show how the metric was aggregated. In this case there is only a single metric so the aggregation is ignored. This is followed by || and then the summary name.
One benefit of this function transformation is it can be nested with other jax function transformations. For example we can jit the transformed function like so:
End of explanation
"""


def layer(params):
  to_look_at = jnp.mean(params) * 2.
  summary.summary("to_look_at", to_look_at)
  return params * 2


@jax.jit
def loss(parameters):
  loss = jnp.mean(layer(layer(parameters))**2)
  return loss


result, metrics = summary.with_summary_output_reduced(loss)(1.)
result, metrics

"""
Explanation: At this point aggregate_metric_list cannot be jit. In practice this is fine as it performs very little computation.
Aggregation of the same name summaries.
Consider the following fake function which calls a layer function twice.
This layer function creates a summary and thus two summary are created.
When running the transformed function we see not one, but two values returned.
End of explanation
"""


summary.aggregate_metric_list([metrics])

"""
Explanation: These values can be combined with aggregate_metric_list as before, but this time the aggregation takes the mean. This mean is specified by the aggregation keyword argument in summary.summary which defaults to mean.
End of explanation
"""


def layer(params):
  to_look_at = jnp.mean(params) * 2.
  summary.summary("to_look_at", to_look_at, aggregation="sample")
  return params * 2


@jax.jit
def loss(parameters):
  loss = jnp.mean(layer(layer(parameters))**2)
  return loss


key = jax.random.PRNGKey(0)
result, metrics = summary.with_summary_output_reduced(loss)(
    1., sample_rng_key=key)
summary.aggregate_metric_list([metrics])

"""
Explanation: Another useful aggregation mode is "sample". As jax's random numbers are stateless, an additional RNG key must be passed in for this to work.
End of explanation
"""


def layer(params):
  to_look_at = jnp.mean(params) * 2.
  summary.summary(
      "to_look_at", jnp.arange(10) * to_look_at, aggregation="collect")
  return params * 2


@jax.jit
def loss(parameters):
  loss = jnp.mean(layer(layer(parameters))**2)
  return loss


key = jax.random.PRNGKey(0)
result, metrics = summary.with_summary_output_reduced(loss)(
    1., sample_rng_key=key)
summary.aggregate_metric_list([metrics])

"""
Explanation: Finally, there is "collect" which concatenates all the values together into one long tensor after first raveling all the inputs. This is useful for extracting distributions of quantities.
End of explanation
"""


@jax.jit
def loss(parameters):
  with summary.summary_scope("scope1"):
    summary.summary("to_look_at", parameters)

  with summary.summary_scope("nested"):
    summary.summary("summary2", parameters)

    with summary.summary_scope("scope2"):
      summary.summary("to_look_at", parameters)
  return parameters


key = jax.random.PRNGKey(0)
result, metrics = summary.with_summary_output_reduced(loss)(
    1., sample_rng_key=key)
summary.aggregate_metric_list([metrics])

"""
Explanation: Summary Scope
Sometimes it is useful to be able to group all summaries from a function inside some code block with some common name. This can be done with the summary_scope context.
End of explanation
"""


@jax.jit
def fn(a):
  summary.summary("other_val", a[2])

  def update(state, _):
    s = state + 1
    summary.summary("mean_loop", s[0])
    summary.summary("collect_loop", s[0], aggregation="collect")
    return s, s

  a, _ = jax.lax.scan(update, a, jnp.arange(20))
  return a * 2


vmap_fn = jax.vmap(fn)

result, metrics = jax.jit(summary.with_summary_output_reduced(vmap_fn))(
    jnp.tile(jnp.arange(4), (2, 2)))
summary.aggregate_metric_list([metrics])

"""
Explanation: Usage with function transforms.
Thanks to oryx, all of this functionality works well across a variety of function transformations.
Here is an example with a scan, vmap, and jit.
The aggregation modes will aggregate across all timesteps, and all batched dimensions.
End of explanation
"""


from learned_optimization import summary
import functools


def layer(params):
  to_look_at = jnp.mean(params) * 2.
  summary.summary("to_look_at", jnp.arange(10) * to_look_at)
  return params * 2


@functools.partial(jax.jit, static_argnames="with_summary")
@summary.add_with_summary
def loss(parameters):
  loss = jnp.mean(layer(layer(parameters))**2)
  return loss


res, metrics = loss(1., with_summary=False)
print("No metrics", summary.aggregate_metric_list([metrics]))

res, metrics = loss(1., with_summary=True)
print("With metrics", summary.aggregate_metric_list([metrics]))

"""
Explanation: Optionally compute metrics: @add_with_metrics
Oftentimes it is useful to define two "versions" of a function -- one with metrics, and one without -- as sometimes the computation of the metrics adds unneeded overhead that does not need to be run every iteration.
To create these two versions one can simply wrap the function with the add_with_summary decorator.
This adds both a keyword argument, and an extra return to the wrapped function which switches between computing metrics, or not.
End of explanation
"""


def monitor(a):
  summary.summary("with_input", a)
  summary.summary("constant", 2.0)
  summary.summary("constant_with_inp", 2.0 + (a * 0))
  return a


result, metrics = summary.with_summary_output_reduced(monitor)(1.)
summary.aggregate_metric_list([metrics])

"""
Explanation: Limitations and Gotchas
Requires value traced to be a function of input
At the moment summary.summary MUST be called with a descendant of whatever input is passed into the function wrapped by with_summary_output_reduced. In practice this is almost always the case as we seek to monitor changing values rather than constants.
To demonstrate this, note how the constant value is NOT logged out, but if add it to a*0 it does become logged out.
End of explanation
"""


def lossb(p):
  to_look_at = jnp.mean(123.)
  return p * 2, to_look_at


def loss(parameters):
  l = jnp.mean(parameters**2)
  l, to_look_at = lossb(l)
  return l, to_look_at


value_grad_fn = jax.jit(jax.value_and_grad(loss, has_aux=True))
(loss, to_look_at), g = value_grad_fn(1.0)
print(to_look_at)

"""
Explanation: The rational for why this is a bit of a rabbit hole, but it is related to how tracing in jax work and is beyond the scope of this notebook.
No support for jax.lax.cond
At this point one cannot extract summaries out of jax conditionals. Sorry. If this is a limitation to you let us know as we have some ideas to make this work.
No dynamic names
At the moment, the tag, or the name of the summary, must be a string known at compile time. There is no support for dynamic summary names.
Alternatives for extracting information
Using this module is not the only way to extract information from a model. We discuss a couple other approaches.
"Thread" metrics through
One way to extract data from a function is to simply return the things we want to look at. As functions become more complex and nested this can become quite a pain as each one of these functions must pass out metric values. This process of spreading data throughout a bunch of functions is called "threading".
Threading also requires all pieces of code to be involved -- as one must thread these metrics everywhere.
End of explanation
"""


from jax.experimental import host_callback as hcb


def loss(parameters):
  loss = jnp.mean(parameters**2)
  to_look_at = jnp.mean(123.)
  hcb.id_print(to_look_at, name="to_look_at")
  return loss


value_grad_fn = jax.jit(jax.value_and_grad(loss))
_ = value_grad_fn(1.0)

"""
Explanation: jax.experimental.host_callback
Jax has some support to send data back from an accelerator back to the host while a ja program is running. This is exposed in jax.experimental.host_callback.
One can use this to print which is a quick way to get data out of a network.
End of explanation
"""


