galax.function

Built-in functions.

Functions

apply_edges(function[, in_field, out_field, ...])	Apply a function to edge attributes.
apply_nodes(function[, in_field, out_field, ...])	Apply a function to node attributes.
copy_e(e, out)	Builtin message function that computes message using edge feature.
copy_edge(e, out)	Builtin message function that computes message using edge feature.
copy_src(u, out)	Builtin message function that computes message using source node feature.
copy_u(u, out)	Builtin message function that computes message using source node feature.
e_add_u(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise add between features of edge and source if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
e_add_v(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise add between features of edge and destination if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
e_div_u(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise div between features of edge and source if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
e_div_v(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise div between features of edge and destination if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
e_dot_u(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise dot between features of edge and source if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
e_dot_v(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise dot between features of edge and destination if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
e_mul_u(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise mul between features of edge and source if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
e_mul_v(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise mul between features of edge and destination if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
e_sub_u(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise sub between features of edge and source if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
e_sub_v(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise sub between features of edge and destination if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
segment_max(*args, **kwargs)	Alias of jax.ops.segment_max with nan_to_num.
segment_mean(data, segment_ids[, ...])	Returns mean for each segment.
segment_min(*args, **kwargs)	Alias of jax.ops.segment_min with nan_to_num.
segment_softmax(data, segment_ids[, ...])	Computes a segment-wise softmax. For a given tree of logits that can be divded into segments, computes a softmax over the segments. logits = jnp.ndarray([1.0, 2.0, 3.0, 1.0, 2.0]) segment_ids = jnp.ndarray([0, 0, 0, 1, 1]) segment_softmax(logits, segments) >> DeviceArray([0.09003057, 0.24472848, 0.66524094, 0.26894142, 0.7310586], >> dtype=float32) Args: logits: an array of logits to be segment softmaxed. segment_ids: an array with integer dtype that indicates the segments of data (along its leading axis) to be maxed over. Values can be repeated and need not be sorted. Values outside of the range [0, num_segments) are dropped and do not contribute to the result. num_segments: optional, an int with positive value indicating the number of segments. The default is jnp.maximum(jnp.max(segment_ids) + 1,   jnp.max(-segment_ids)) but since num_segments determines the size of the output, a static value must be provided to use segment_sum in a jit-compiled function. indices_are_sorted: whether segment_ids is known to be sorted unique_indices: whether segment_ids is known to be free of duplicates Returns: The segment softmax-ed logits.
u_add_e(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise add between features of source and edge if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
u_add_v(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise add between features of source and destination if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
u_div_e(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise div between features of source and edge if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
u_div_v(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise div between features of source and destination if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
u_dot_e(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise dot between features of source and edge if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
u_dot_v(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise dot between features of source and destination if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
u_mul_e(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise mul between features of source and edge if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
u_mul_v(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise mul between features of source and destination if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
u_sub_e(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise sub between features of source and edge if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
u_sub_v(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise sub between features of source and destination if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
v_add_e(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise add between features of destination and edge if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
v_add_u(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise add between features of destination and source if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
v_div_e(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise div between features of destination and edge if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
v_div_u(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise div between features of destination and source if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
v_dot_e(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise dot between features of destination and edge if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
v_dot_u(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise dot between features of destination and source if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
v_mul_e(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise mul between features of destination and edge if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
v_mul_u(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise mul between features of destination and source if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
v_sub_e(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise sub between features of destination and edge if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.
v_sub_u(lhs_field, rhs_field, out)	Builtin message function that computes a message on an edge by performing element-wise sub between features of destination and source if the features have the same shape; otherwise, it first broadcasts the features to a new shape and performs the element-wise operation.

Classes

ReduceFunction(op, msg_field, out_field)