class JaxLayer:
    pass