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.opsfor 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.KerasSalienceModelto 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.