LIT: Refactor instrumented Keras LMs to support TF and Torch
Summary of changes:
Wherever possible, use keras.ops for Tensor operations. These functions are designed to create and operate over Tensors from any of the Keras 3 backends.
Identified a common API surface for salience predictions that supports all existing (and likely any future) Keras 3 backends. Functions implementing this API shoudl be named _pred_{framework} and return a tuple of the GradNorm and GradDotInput salience scores.
Refactors KerasSalienceModel._pred() to perform common operations for data before calling out to backend-specific prediction functions.
Extracts TensorFlow code in KerasSalienceModel to a new _pred_tf() function.
Implements a KerasSalienceModel._pred_torch() function based on the HF implementation in lit_nlp/examples/models/pretrained_lms.py
Provides a stub for KerasSalienceModel._pred_jax() with a detailed comment outlining the JAX idiosyncrasies that we need to adapt to in order to support this backend.
LIT: Refactor instrumented Keras LMs to support TF and Torch
Summary of changes:
keras.ops
for Tensor operations. These functions are designed to create and operate over Tensors from any of the Keras 3 backends._pred_{framework}
and return a tuple of the GradNorm and GradDotInput salience scores.KerasSalienceModel._pred()
to perform common operations for data before calling out to backend-specific prediction functions.KerasSalienceModel
to a new_pred_tf()
function.KerasSalienceModel._pred_torch()
function based on the HF implementation in lit_nlp/examples/models/pretrained_lms.pyKerasSalienceModel._pred_jax()
with a detailed comment outlining the JAX idiosyncrasies that we need to adapt to in order to support this backend.