flax

4183e15c - Copybara import of the project:

Copybara import of the project: -- 5074fb68b3d61f9ad47d26d61435298ea767573e by vfdev-5 <vfdev.5@gmail.com>: Added support for data masking in Average, Accuracy and MultiMetric PiperOrigin-RevId: 883351273