JIT Compile a Function

Compile a function that applies sin(x)**2 + cos(x)**2 to every tensor value. This exercise practices compilation, so keep the trigonometric operations instead of replacing them with a tensor of ones.

Your task

Implement jit_trig_identity(x) using @jax.jit in JAX or @torch.compile in PyTorch. The input is a floating-point tensor of finite angles in radians. Return a tensor with the same shape, dtype and device, leaving x unchanged.

The identity sin(t)**2 + cos(t)**2 = 1 holds for real numbers in exact arithmetic. Computed values should be close to one, with small differences from floating-point rounding. For input [0, 1, 2], the expected result is approximately [1, 1, 1]. Do not return a single scalar or reduce over an axis.

Define the function with framework operations, then decorate it or create a compiled wrapper once and reuse it. Return the computed tensor, not the wrapper. Python uses ** for powers; ^ is not exponentiation.

What compilation changes

The first call may include tracing and compilation work. Later calls can reuse compiled work when their inputs satisfy the compiler's requirements. Compilation does not guarantee a speedup for this tiny function, especially if the timing includes the first call.

The diagram is a simplified JAX cache model for one function. The first float32 vector of length 32 creates an entry; another call of the same shape and dtype reuses it; a vector of length 64 creates another entry. The labels refer to vector lengths, not float32 versus float64 precision.

Real cache behavior also depends on the function and compilation settings. PyTorch uses guards and can support dynamic shapes, so this diagram does not predict its exact recompilation behavior. It simulates cache decisions; it does not compile or time code in your browser.