hgq.layers.attn package

Submodules

hgq.layers.attn.linformer module

class hgq.layers.attn.linformer.QLinformerAttention(*args, **kwargs)

Bases: QMultiHeadAttention

build(query_shape, value_shape=None, key_shape=None)

Builds layers and variables.

Parameters:

  • query_shape (tuple) – Shape of the query tensor.

  • value_shape (tuple, optional) – Shape of the value tensor.

  • key_shape (tuple, optional) – Shape of the key tensor.

call(query, value=None, key=None, query_mask=None, value_mask=None, key_mask=None, attention_mask=None, return_attention_scores=False, training=None, use_causal_mask=False)
property ebops
get_config()

Returns the config of the object.

An object config is a Python dictionary (serializable) containing the information needed to re-instantiate it.

hgq.layers.attn.mha module

class hgq.layers.attn.mha.QMultiHeadAttention(*args, **kwargs)

Bases: MultiHeadAttention, QLayerBase

build(query_shape, value_shape=None, key_shape=None)

Builds layers and variables.

Parameters:

  • query_shape (tuple) – Shape of the query tensor.

  • value_shape (tuple, optional) – Shape of the value tensor.

  • key_shape (tuple, optional) – Shape of the key tensor.

call(query, value=None, key=None, query_mask=None, value_mask=None, key_mask=None, attention_mask=None, return_attention_scores=False, training=None, use_causal_mask=False)
compute_output_shape(query_shape, value_shape, key_shape=None)
compute_output_spec(query, value=None, key=None, *args, **kwargs)
property ebops
get_config()

Returns the config of the object.

An object config is a Python dictionary (serializable) containing the information needed to re-instantiate it.

hgq.layers.attn.salt module

class hgq.layers.attn.salt.QSALTAttention(*args, **kwargs)

Bases: QMultiHeadAttention

Quantized version of SALT attention from “Spatially Aware Linear Transformer (SALT) for Particle Jet Tagging” (https://arxiv.org/abs/2510.23641). Adds a conv2d operation on the computed attnmask before applying it to the value tensor.

build(query_shape, value_shape=None, key_shape=None)

Builds layers and variables.

Parameters:

  • query_shape (tuple) – Shape of the query tensor.

  • value_shape (tuple, optional) – Shape of the value tensor.

  • key_shape (tuple, optional) – Shape of the key tensor.

call(query, value=None, key=None, query_mask=None, value_mask=None, key_mask=None, attention_mask=None, return_attention_scores=False, training=None, use_causal_mask=False)
property ebops
get_config()

Returns the config of the object.

An object config is a Python dictionary (serializable) containing the information needed to re-instantiate it.

Module contents

class hgq.layers.attn.QLinformerAttention(*args, **kwargs)

Bases: QMultiHeadAttention

build(query_shape, value_shape=None, key_shape=None)

Builds layers and variables.

Parameters:

  • query_shape (tuple) – Shape of the query tensor.

  • value_shape (tuple, optional) – Shape of the value tensor.

  • key_shape (tuple, optional) – Shape of the key tensor.

call(query, value=None, key=None, query_mask=None, value_mask=None, key_mask=None, attention_mask=None, return_attention_scores=False, training=None, use_causal_mask=False)
property ebops
get_config()

Returns the config of the object.

An object config is a Python dictionary (serializable) containing the information needed to re-instantiate it.

class hgq.layers.attn.QMultiHeadAttention(*args, **kwargs)

Bases: MultiHeadAttention, QLayerBase

build(query_shape, value_shape=None, key_shape=None)

Builds layers and variables.

Parameters:

  • query_shape (tuple) – Shape of the query tensor.

  • value_shape (tuple, optional) – Shape of the value tensor.

  • key_shape (tuple, optional) – Shape of the key tensor.

call(query, value=None, key=None, query_mask=None, value_mask=None, key_mask=None, attention_mask=None, return_attention_scores=False, training=None, use_causal_mask=False)
compute_output_shape(query_shape, value_shape, key_shape=None)
compute_output_spec(query, value=None, key=None, *args, **kwargs)
property ebops
get_config()

Returns the config of the object.

An object config is a Python dictionary (serializable) containing the information needed to re-instantiate it.

class hgq.layers.attn.QSALTAttention(*args, **kwargs)

Bases: QMultiHeadAttention

Quantized version of SALT attention from “Spatially Aware Linear Transformer (SALT) for Particle Jet Tagging” (https://arxiv.org/abs/2510.23641). Adds a conv2d operation on the computed attnmask before applying it to the value tensor.

build(query_shape, value_shape=None, key_shape=None)

Builds layers and variables.

Parameters:

  • query_shape (tuple) – Shape of the query tensor.

  • value_shape (tuple, optional) – Shape of the value tensor.

  • key_shape (tuple, optional) – Shape of the key tensor.

call(query, value=None, key=None, query_mask=None, value_mask=None, key_mask=None, attention_mask=None, return_attention_scores=False, training=None, use_causal_mask=False)
property ebops
get_config()

Returns the config of the object.

An object config is a Python dictionary (serializable) containing the information needed to re-instantiate it.