FR: Latex rendering support for treescope axis names

[amifalk]

e.g. treescope.render_array(pz.nx.ones({r"$\omega$": 5}))

[amifalk]

Here's a hack that allows rendering equations in treescope with MathJax. It seems like it should be possible to inject the MathJax script somewhere in the treescope rendering code so it applies to the text in all figures, but I haven't been able to find the right spot to put it.

import treescope as ts
from penzai import pz

ts.basic_interactive_setup()


class MathJax:
    def __init__(self, value):
        self.value = value

    def _repr_html_(self):
        mathjax_script = (
            '<script id="MathJax-script" async'
            ' src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js">'
            ' </script>"'
        )

        return mathjax_script + self.value


equation = MathJax(r"\(f(x)=\int{x^2}\), lorem ipsum")
arr = pz.nx.ones({"a": 10, "b": 10})

ts.show(equation, "\n treescope text\n", arr)